diff --git a/arrival_times.py b/arrival_times.py index aee164b..97e5c58 100644 --- a/arrival_times.py +++ b/arrival_times.py @@ -3,9 +3,10 @@ import datetime class ArrivalTime(): """ Represents the arrival times of buses at one of the configured stops """ - def __init__(self, stop_id: str, route_id: str, destination: str, due_in_seconds: int, is_added: bool = False) -> None: + def __init__(self, stop_id: str, route_id: str, route_type: str, destination: str, due_in_seconds: int, is_added: bool = False) -> None: self.stop_id = stop_id self.route_id = route_id + self.route_type = route_type self.destination = destination self.due_in_seconds = due_in_seconds self.is_added = is_added diff --git a/gtfs_client.py b/gtfs_client.py index 2eedb63..d3cc60c 100644 --- a/gtfs_client.py +++ b/gtfs_client.py @@ -316,6 +316,7 @@ class GTFSClient(): new_arrival = ArrivalTime( stop_id = stop_time_update.get("stop_code"), route_id = self.feed.routes[self.feed.routes["route_id"] == route_id]["route_short_name"].item(), + route_type = self.feed.routes[self.feed.routes["route_id"] == route_id]["route_type"].item(), destination = self.__lookup_headsign_by_route(route_id, direction_id), due_in_seconds = arrival_time - int(time.time()), is_added = True @@ -376,6 +377,7 @@ class GTFSClient(): arrival = ArrivalTime(stop_id = bus["stop_code"], route_id = bus["route_short_name"], + route_type = bus["route_type"], destination = bus["trip_headsign"], due_in_seconds = self.__due_in_seconds(bus["arrival_time"]) + delta, is_added = False