43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
import os
|
|
import sys
|
|
import zipfile
|
|
import pandas as pd
|
|
import tempfile
|
|
|
|
|
|
def modify_routes(gtfs_file, routes, output_file=None):
|
|
if not output_file:
|
|
output_file = gtfs_file
|
|
|
|
# unzip gtfs feed
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
with zipfile.ZipFile(gtfs_file, 'r') as zip_ref:
|
|
zip_ref.extractall(tmpdirname)
|
|
|
|
# load routes.txt in a DataFrame
|
|
routes_path = os.path.join(tmpdirname, 'routes.txt')
|
|
routes_df = pd.read_csv(routes_path)
|
|
|
|
# edit route types
|
|
for route_change in routes.split(','):
|
|
route_id, new_type = route_change.split(':')
|
|
routes_df.loc[routes_df['route_id'] == route_id, 'route_type'] = int(new_type)
|
|
|
|
# save to routes.txt
|
|
routes_df.to_csv(routes_path, index=False)
|
|
|
|
# recreate zip file
|
|
with zipfile.ZipFile(output_file, 'w') as zip_ref:
|
|
for foldername, subfolders, filenames in os.walk(tmpdirname):
|
|
for filename in filenames:
|
|
file_path = os.path.join(foldername, filename)
|
|
zip_ref.write(file_path, os.path.relpath(file_path, tmpdirname))
|
|
|
|
|
|
if __name__ == "main":
|
|
gtfs_file = sys.argv[1]
|
|
routes = sys.argv[2]
|
|
output_file = sys.argv[3]
|
|
|
|
modify_routes(gtfs_file, routes, output_file)
|