make_stop_map.py 1.32 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49

import csv
import json


def reformat_stop(stop):
    return {
        "n": stop["stop_name"],
        "x": float(stop["stop_lon"]),
        "y": float(stop["stop_lat"]),
        "r": stop["routes"]
    }


def reformat_route(route):
    return {
        "n": route["route_long_name"]
    }


def main():
    routes = {}
    stops = {}
    trips = {}
    with open('google_bus/routes.txt', 'r') as f:
        for route in csv.DictReader(f):
            routes[route["route_id"]] = route
    with open('google_bus/stops.txt', 'r') as f:
        for stop in csv.DictReader(f):
            stops[stop["stop_id"]] = stop
            stops[stop["stop_id"]]["routes"] = []
    with open('google_bus/trips.txt', 'r') as f:
        for trip in csv.DictReader(f):
            trips[trip["trip_id"]] = trip
    with open('google_bus/stop_times.txt', 'r') as f:
        for stop_time in csv.DictReader(f):
            route_id = trips[stop_time["trip_id"]]["route_id"]
            stop = stops[stop_time["stop_id"]]
            if route_id not in stop["routes"]:
                stop["routes"].append(route_id)
    output = {
        "routes": dict([(k, reformat_route(route)) for k, route in routes.items()]),
        "stops": [reformat_stop(stop) for stop in stops.values()]
    }
    with open("stops.json", "w") as f:
        json.dump(output, f)


main()