RideLogicAVLS/experimental/geojson_from_gtfs_txt.py

#!/usr/bin/python3
#
#
# Copyright (c) 2021 Clementine Computing LLC.
# 
# This file is part of RideLogic.
# 
# RideLogic is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# RideLogic is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
# 
# You should have received a copy of the GNU Affero General Public License
# along with RideLogic.  If not, see <https://www.gnu.org/licenses/>.
#

import os
import sys
import csv
import json

if len(sys.argv) < 6:
  print("provide routes.txt stops.txt shapes.txt trips.txt stop_times.txt")
  sys.exit(-1)

routes_fn = sys.argv[1]
stops_fn = sys.argv[2]
shapes_fn = sys.argv[3]
trips_fn = sys.argv[4]
stop_times_fn = sys.argv[5]

route_info = {}
stop_info = {}
shape_info = {}
trip_info = {}
stop_time_info = {}

with open(routes_fn) as fp:
  fpcsv = csv.reader(fp, delimiter=',', quotechar='"')
  next(fpcsv)
  for row in fpcsv:
    route_info[row[0]] = {
      "route_id" : row[0],
      "agency_id" : row[1],
      "route_short_name": row[2],
      "route_long_name" : row[3],
      "route_desc" : row[4],
      "route_type": row[5],
      "route_url" : row[6],
      "route_color" : row[7],
      "route_text_color" : row[8],
      "route_sort_order" : row[9]
    }

with open(stops_fn) as fp:
  fpcsv = csv.reader(fp, delimiter=',', quotechar='"')
  next(fpcsv)
  for row in fpcsv:
    stop_info[row[0]] = {
      "stop_id" : row[0],
      "stop_code" : row[1],
      "stop_name": row[2],
      "stop_desc" : row[3],
      "stop_lat" : row[4],
      "stop_lon": row[5],
      "zone_id" : row[6],
      "stop_url" : row[7],
      "location_type" : row[8],
      "parent_station" : row[9],
      "wheelchair_boarding" : row[10]
    }

with open(shapes_fn) as fp:
  fpcsv = csv.reader(fp, delimiter=',', quotechar='"')
  next(fpcsv)
  for row in fpcsv:
    if not (row[0] in shape_info):
      shape_info[row[0]] = []
    shape_info[row[0]].append({
      "shape_id" : row[0],
      "shape_pt_lat" : row[1],
      "shape_pt_lon": row[2],
      "shape_pt_sequence" : row[3],
      "shape_dist_traveled" : row[4]
    })

with open(trips_fn) as fp:
  fpcsv = csv.reader(fp, delimiter=',', quotechar='"')
  next(fpcsv)
  for row in fpcsv:
    trip_info[row[2]] = {
      "route_id" : row[0],
      "service_id": row[1],
      "trip_id" : row[2],
      "trip_headsign" : row[3],
      "direction_id" : row[4],
      "block_id": row[5],
      "shape_id" : row[6],
      "wheelchair_accessible" : row[7],
      "bikes_allowed" : row[8]
    }

trip_stops_map = {}

with open(stop_times_fn) as fp:
  fpcsv = csv.reader(fp, delimiter=',', quotechar='"')
  next(fpcsv)
  for row in fpcsv:
    trip_id = row[0]
    stop_id = row[3]
    if not (trip_id in stop_time_info):
      stop_time_info[trip_id] = []
    stop_time_info[trip_id].append({
      "trip_id" : row[0],
      "arrival_time" : row[1],
      "departure_time": row[2],
      "stop_id" : row[3],
      "stop_sequence" : row[4],
      "stop_headsign" : row[5],
      "pickup_type" : row[6],
      "srop_off_type" : row[7],
      "shape_dist_traveled" : row[8],
      "timepoint" : row[9]
    })


    if not (trip_id in trip_stops_map):
      trip_stops_map[trip_id] = {}

    if not (stop_id in trip_stops_map[trip_id]):
      trip_stops_map[trip_id][stop_id] = True


geojson_info = {
  "type" : "FeatureCollection",
  "features" : []
}

for trip_id in trip_info:
  geojson_trip = {
    "type" : "trip",
    "geometry": {
      "type": "Polygon",
      "coordinates": []
    },
    "properties" : {}
  }
  trip = trip_info[trip_id]
  shape_id = trip["shape_id"]

  for shape_ele in shape_info[shape_id]:
    #geojson_trip["geometry"]["coordinates"].append( [ shape_ele["shape_pt_lat"], shape_ele["shape_pt_lon"] ] )
    geojson_trip["geometry"]["coordinates"].append( [ shape_ele["shape_pt_lon"], shape_ele["shape_pt_lat"] ] )

  #geojson_trip["properties"]["route_id"] = trip["route_id"]
  geojson_trip["properties"]["trip_id"] = trip["trip_id"]
  geojson_trip["properties"]["service_id"] = trip["service_id"]
  geojson_trip["properties"]["block_id"] = trip["block_id"]
  geojson_trip["properties"]["trip_headsign"] = trip["trip_headsign"]
  geojson_trip["properties"]["bikes_allowed"] = trip["bikes_allowed"]
  geojson_trip["properties"]["wheelchair_accessible"] = trip["wheelchair_accessible"]

  route_ele = route_info[ trip["route_id"] ]

  route_stops = []
  for stop_id in trip_stops_map[trip_id]:
    stop_ele = stop_info[stop_id]

    gj_stop_ele = {
      "stop" : {
        "geometry": {
          "coordinates" : [stop_ele["stop_lon"], stop_ele["stop_lat"] ]
        }
      }
    }
    route_stops.append( gj_stop_ele )


  geojson_trip["properties"]["route"] = {
    "route_id" : route_ele["route_id"],
    "agency_id" : route_ele["agency_id"],
    "route_short_name" : route_ele["route_short_name"],
    "route_long_name" : route_ele["route_long_name"],
    "route_desc" : route_ele["route_desc"],
    "route_type" : route_ele["route_type"],
    "route_url" : route_ele["route_url"],
    "route_color" : route_ele["route_color"],
    "route_text_color" : route_ele["route_text_color"],
    "route_sort_order" : route_ele["route_sort_order"],
  }
  geojson_trip["properties"]["route_stops"] = route_stops

  geojson_info["features"].append(geojson_trip)

  #print( json.dumps(geojson_trip, indent=2) )
  #print( json.dumps(trip_info[trip_id], indent=2) )

#print(json.dumps(route_info, indent=2))
#print(json.dumps(stop_info, indent=2))
#print(json.dumps(shape_info, indent=2))
#print(json.dumps(trip_info, indent=2))

print(json.dumps(geojson_info, indent=2))