change-route-type/modify_gtfs.py
Charles P. 83193a4adb
Some checks are pending
GTFS Route Modification / modify_routes (push) Waiting to run
🚧 add arguments to dockerfile
2024-05-31 22:11:49 +02:00

58 lines
1.9 KiB
Python

import os
import sys
import zipfile
import pandas as pd
import tempfile
import logging
def modify_routes(gtfs_file, routes, output_file=None):
logging.info(f"Running GTFS Route Type Modifier")
logging.info(f"GTFS File: {gtfs_file}")
logging.info(f"Routes to Change: {routes}")
logging.info(f"Output File: {output_file}")
if not output_file:
output_file = gtfs_file
# unzip gtfs feed
with tempfile.TemporaryDirectory() as tmpdirname:
logging.info(f"Unzipping GTFS file to {tmpdirname}")
with zipfile.ZipFile(gtfs_file, 'r') as zip_ref:
zip_ref.extractall(tmpdirname)
# load routes.txt in a DataFrame
logging.info(f"Reading routes.txt")
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(':')
logging.info(f"Changing route {route_id} to type {new_type}")
routes_df.loc[routes_df['route_id'] == route_id, 'route_type'] = int(new_type)
# save to routes.txt
logging.info(f"Saving routes.txt")
routes_df.to_csv(routes_path, index=False)
# recreate zip file
logging.info(f"Recreating GTFS file at {output_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))
logging.info(f"Done!")
if __name__ == "main":
gtfs_file = sys.argv[1]
routes = sys.argv[2]
output_file = sys.argv[3]
logging.basicConfig(level=logging.INFO)
modify_routes(gtfs_file, routes, output_file)