diff --git a/modify_gtfs.py b/modify_gtfs.py index 5edc79f..8f1d7bc 100644 --- a/modify_gtfs.py +++ b/modify_gtfs.py @@ -3,40 +3,54 @@ 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)