gtfs.py

import datetime
import csv
from dataclasses import dataclass
import numpy as np
import os, sys
from utils import *
import weakref

local_timezone = datetime.datetime.now(datetime.timezone.utc).astimezone().tzinfo

_shape_loading_from_disk_count = 0
_shape_loading_from_cache_count = 0

@dataclass
class Route:
    id: str
    name: str

@dataclass
class Trip:
    route: Route
    id: int
    direction: int
    shape_id: str
    services_today: bool


@dataclass
class Stop:
    id: int
    name: str
    lat: float
    lon: float


@dataclass
class ShapePoint:
    lat: str
    lon: float
    dist_traveled: float

@dataclass
class TripStop:
    stop: Stop
    arrival_time: float
    departure_time: float
    shape_dist_traveled: float


class GtfsDay:
    def __init__(self, date, data_getter):
        self.date = date
        self.data_getter = data_getter
        self.stops = None
        self.trips = None
        self.trips_by_routes = None
        self.routes = None
        self.services_today = None
        self.stops_for_trip = None
        self._shape_cache = weakref.WeakValueDictionary()
        self._file_cahce = {}

    async def get_file(self, name, cache=False):
        if name in self._file_cahce:
            return self._file_cahce[name]
        s = await self.data_getter(datetime.datetime.combine(self.date, datetime.time()), name) # HACK: Date is not CBOR supported type
        r =  list(csv.DictReader(s.decode("utf-8").split("\n")))
        if cache:
            self._file_cahce[name] = r
        return r

    async def load_calendar(self):
        if self.services_today: return
        date = self.date.strftime('%Y%m%d')
        d = await self.get_file("calendar.txt")
        assert len(d) == len({x['service_id'] for x in d}), "Duplicit service_id"
        services_today = {x['service_id']: x[['monday','tuesday','wednesday','thursday','friday','saturday','sunday'][self.date.weekday()]] == '1' and x['start_date']<=date  and date<=x['end_date'] for x in d}

        d = await self.get_file("calendar_dates.txt")
        for x in d:
            if date == x['date']:
                services_today[x['service_id']] = x['exception_type'] == 1
        self.services_today = services_today

    async def load_trips(self):
        await self.load_calendar()
        if self.trips: return

        d = await self.get_file("routes.txt")
        self.routes = {
                x["route_id"]: Route(
                    x["route_id"],
                    x["route_short_name"],
                ) for x in d
        }

        d = await self.get_file("trips.txt")
        self.trips = {
                x["trip_id"]: Trip(
                    route=self.routes[x["route_id"]],
                    id=x["trip_id"],
                    direction=int(x["direction_id"]),
                    shape_id=x["shape_id"],
                    services_today=self.services_today[x["service_id"]]
                ) for x in d
        }
        self.trips_by_routes = {}
        for t in self.trips.values():
            r_id = t.route.id
            self.trips_by_routes.setdefault(r_id, [])
            self.trips_by_routes[r_id].append(t)


    async def load_stops(self):
        if self.stops: return
        d = await self.get_file("stops.txt")
        if self.stops: return
        self.stops = {
                x["stop_id"]: Stop(x["stop_id"], x["stop_name"], float(x["stop_lat"]), float(x["stop_lon"]))
                for x in d
        }

    async def get_shape_for_trip_id(self, trip_id):
        global _shape_loading_from_disk_count, _shape_loading_from_cache_count
        await self.load_trips()
        if trip_id not in self.trips:
            eprint(f"get_shape_for_trip_id failed: no such trip_id {trip_id}")
            return None
        shape_id = self.trips[trip_id].shape_id
        shape = self._shape_cache.get(shape_id, None)
        if shape is None:
            _shape_loading_from_disk_count += 1
            d = await self.get_file("shape_by_id/"+shape_id)
            self._shape_cache[shape_id] = shape = np.array([
                [float(x["shape_pt_lat"]), float(x["shape_pt_lon"]), float(x["shape_dist_traveled"])]

            for x in d ])
        else:
            _shape_loading_from_cache_count += 1
        return shape

    def parse_time(self, val):
        h, m, s  = map(int, val.split(":"))
        return datetime.datetime.combine(self.date, datetime.time(0, 0, 0), local_timezone) + datetime.timedelta(hours=h, minutes=m, seconds=s) # hack for h > 23

    async def get_stops_for_trip_id(self, trip_id, data=None):
        await self.load_trips()
        await self.load_stops()
        if self.stops_for_trip:
            if not trip_id in self.stops_for_trip:
                eprint(f"get_stops_for_trip_id failed: no such trip {trip_id} at {self.date}")
                return []
            return self.stops_for_trip[trip_id]
        if data is not None:
            d =  list(csv.DictReader(data.decode("utf-8").split("\n")))
        else:
            eprint("LOADING STOPS FOR", trip_id, self.date, "(without cache)")
            d = await self.get_file("stop_times.txt", cache=False)
        return [ TripStop(
            self.stops[x["stop_id"]],
            self.parse_time(x["arrival_time"]),
            self.parse_time(x["departure_time"]),
            float(x["shape_dist_traveled"])
        ) for x in d if x["trip_id"] == trip_id]

    async def load_stops_for_all_trips(self):
        if self.stops_for_trip:
            return
        await self.load_trips()
        await self.load_stops()
        d = await self.get_file("stop_times.txt")
        stops_for_trip = {}
        for x in d:
            trip_id = x["trip_id"]
            if not trip_id in stops_for_trip:
                stops_for_trip[trip_id] = []
            stops_for_trip[trip_id].append(
                TripStop(
                    self.stops[x["stop_id"]],
                    self.parse_time(x["arrival_time"]),
                    self.parse_time(x["departure_time"]),
                    float(x["shape_dist_traveled"])
                )
            )

        self.stops_for_trip = stops_for_trip







default_data_getter = None

for_date_cache = {}

def for_date(date, data_getter=None):
    if isinstance(date, datetime.datetime):
        date = date.date()
    if date not in for_date_cache:
        for_date_cache[date] = GtfsDay(date, data_getter or default_data_getter)
    return for_date_cache[date]