change-route-type/modify_gtfs.py

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 Change Route Type")
    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)