From 175c2ce99a4aa8e35c19d0e1d9845e8ecfd35814 Mon Sep 17 00:00:00 2001 From: Jason Goode <53153201+jasonmgoode@users.noreply.github.com> Date: Tue, 28 Sep 2021 12:46:52 -0400 Subject: [PATCH] Add MBTA translator (#38) --- .../registry/registry.py | 3 +- .../translators/__init__.py | 1 + gtfs_realtime_translators/translators/mbta.py | 49 +++ test/fixtures/mbta_bus.json | 308 ++++++++++++++ test/fixtures/mbta_subway.json | 387 ++++++++++++++++++ test/test_mbta.py | 54 +++ 6 files changed, 801 insertions(+), 1 deletion(-) create mode 100644 gtfs_realtime_translators/translators/mbta.py create mode 100644 test/fixtures/mbta_bus.json create mode 100644 test/fixtures/mbta_subway.json create mode 100644 test/test_mbta.py diff --git a/gtfs_realtime_translators/registry/registry.py b/gtfs_realtime_translators/registry/registry.py index 312b38d..0377e43 100644 --- a/gtfs_realtime_translators/registry/registry.py +++ b/gtfs_realtime_translators/registry/registry.py @@ -3,7 +3,7 @@ from gtfs_realtime_translators.translators import LaMetroGtfsRealtimeTranslator, \ SeptaRegionalRailTranslator, MtaSubwayGtfsRealtimeTranslator, NjtRailGtfsRealtimeTranslator, \ CtaSubwayGtfsRealtimeTranslator, CtaBusGtfsRealtimeTranslator, PathGtfsRealtimeTranslator, \ - SwiftlyGtfsRealtimeTranslator, WcdotGtfsRealTimeTranslator, NjtBusGtfsRealtimeTranslator + SwiftlyGtfsRealtimeTranslator, WcdotGtfsRealTimeTranslator, NjtBusGtfsRealtimeTranslator, MbtaGtfsRealtimeTranslator class TranslatorKeyWarning(Warning): @@ -23,6 +23,7 @@ class TranslatorRegistry: 'path': PathGtfsRealtimeTranslator, 'vta': SwiftlyGtfsRealtimeTranslator, 'wcdot-bus': WcdotGtfsRealTimeTranslator, + 'mbta': MbtaGtfsRealtimeTranslator, } @classmethod diff --git a/gtfs_realtime_translators/translators/__init__.py b/gtfs_realtime_translators/translators/__init__.py index 66bb524..c03bcc5 100644 --- a/gtfs_realtime_translators/translators/__init__.py +++ b/gtfs_realtime_translators/translators/__init__.py @@ -8,3 +8,4 @@ from .path_rail import PathGtfsRealtimeTranslator from .swiftly import SwiftlyGtfsRealtimeTranslator from .wcdot_bus import WcdotGtfsRealTimeTranslator +from .mbta import MbtaGtfsRealtimeTranslator diff --git a/gtfs_realtime_translators/translators/mbta.py b/gtfs_realtime_translators/translators/mbta.py new file mode 100644 index 0000000..6c5455f --- /dev/null +++ b/gtfs_realtime_translators/translators/mbta.py @@ -0,0 +1,49 @@ +import json + +import pendulum + +from gtfs_realtime_translators.factories import TripUpdate, FeedMessage + + +class MbtaGtfsRealtimeTranslator: + TIMEZONE = 'America/New_York' + + def __call__(self, data): + json_data = json.loads(data) + predictions = json_data['data'] + entities = self.__make_trip_updates(predictions) + return FeedMessage.create(entities=entities) + + + @classmethod + def __to_unix_time(cls, time): + return pendulum.parse(time).in_tz(cls.TIMEZONE).int_timestamp + + + @classmethod + def __make_trip_updates(cls, predictions): + trip_updates = [] + for idx, prediction in enumerate(predictions): + entity_id = str(idx + 1) + relationships = prediction['relationships'] + attributes = prediction['attributes'] + stop_id = relationships['stop']['data']['id'] + route_id = relationships['route']['data']['id'] + trip_id = relationships['trip']['data']['id'] + raw_arrival_time = attributes['arrival_time'] + raw_departure_time = attributes['departure_time'] + + if raw_arrival_time and raw_departure_time: + arrival_time = cls.__to_unix_time(attributes['arrival_time']) + departure_time = cls.__to_unix_time(attributes['departure_time']) + trip_update = TripUpdate.create( + entity_id=entity_id, + route_id=route_id, + stop_id=stop_id, + trip_id=trip_id, + arrival_time=arrival_time, + departure_time=departure_time + ) + trip_updates.append(trip_update) + + return trip_updates diff --git a/test/fixtures/mbta_bus.json b/test/fixtures/mbta_bus.json new file mode 100644 index 0000000..e4d8613 --- /dev/null +++ b/test/fixtures/mbta_bus.json @@ -0,0 +1,308 @@ +{ + "data": [ + { + "attributes": { + "arrival_time": null, + "departure_time": null, + "direction_id": 0, + "schedule_relationship": "SKIPPED", + "status": null, + "stop_sequence": 6 + }, + "id": "prediction-49181402-1357-6", + "relationships": { + "route": { + "data": { + "id": "66", + "type": "route" + } + }, + "stop": { + "data": { + "id": "1357", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49181402", + "type": "trip" + } + }, + "vehicle": { + "data": null + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": null, + "departure_time": null, + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 6 + }, + "id": "prediction-49181421-1357-6", + "relationships": { + "route": { + "data": { + "id": "66", + "type": "route" + } + }, + "stop": { + "data": { + "id": "1357", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49181421", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "y1799", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T17:38:53-04:00", + "departure_time": "2021-09-27T17:38:53-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 6 + }, + "id": "prediction-49181349-1357-6", + "relationships": { + "route": { + "data": { + "id": "66", + "type": "route" + } + }, + "stop": { + "data": { + "id": "1357", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49181349", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "y1885", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T17:51:18-04:00", + "departure_time": "2021-09-27T17:51:18-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 6 + }, + "id": "prediction-49181350-1357-6", + "relationships": { + "route": { + "data": { + "id": "66", + "type": "route" + } + }, + "stop": { + "data": { + "id": "1357", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49181350", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "y1838", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T18:00:15-04:00", + "departure_time": "2021-09-27T18:00:15-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 6 + }, + "id": "prediction-49181351-1357-6", + "relationships": { + "route": { + "data": { + "id": "66", + "type": "route" + } + }, + "stop": { + "data": { + "id": "1357", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49181351", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "y1910", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T18:07:39-04:00", + "departure_time": "2021-09-27T18:07:39-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 6 + }, + "id": "prediction-49181352-1357-6", + "relationships": { + "route": { + "data": { + "id": "66", + "type": "route" + } + }, + "stop": { + "data": { + "id": "1357", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49181352", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "y1792", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T18:18:56-04:00", + "departure_time": "2021-09-27T18:18:56-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 6 + }, + "id": "prediction-49181353-1357-6", + "relationships": { + "route": { + "data": { + "id": "66", + "type": "route" + } + }, + "stop": { + "data": { + "id": "1357", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49181353", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "y1729", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T18:40:03-04:00", + "departure_time": "2021-09-27T18:40:03-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 6 + }, + "id": "prediction-49181355-1357-6", + "relationships": { + "route": { + "data": { + "id": "66", + "type": "route" + } + }, + "stop": { + "data": { + "id": "1357", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49181355", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "y1754", + "type": "vehicle" + } + } + }, + "type": "prediction" + } + ], + "jsonapi": { + "version": "1.0" + } +} \ No newline at end of file diff --git a/test/fixtures/mbta_subway.json b/test/fixtures/mbta_subway.json new file mode 100644 index 0000000..68b2b74 --- /dev/null +++ b/test/fixtures/mbta_subway.json @@ -0,0 +1,387 @@ +{ + "data": [ + { + "attributes": { + "arrival_time": "2021-09-27T15:26:30-04:00", + "departure_time": "2021-09-27T15:27:26-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 70 + }, + "id": "prediction-49417961-70045-70", + "relationships": { + "route": { + "data": { + "id": "Blue", + "type": "route" + } + }, + "stop": { + "data": { + "id": "70045", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49417961", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "B-546C7355", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T15:31:46-04:00", + "departure_time": "2021-09-27T15:32:42-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 70 + }, + "id": "prediction-49417962-70045-70", + "relationships": { + "route": { + "data": { + "id": "Blue", + "type": "route" + } + }, + "stop": { + "data": { + "id": "70045", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49417962", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "B-546C7333", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T15:35:49-04:00", + "departure_time": "2021-09-27T15:36:45-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 70 + }, + "id": "prediction-49417963-70045-70", + "relationships": { + "route": { + "data": { + "id": "Blue", + "type": "route" + } + }, + "stop": { + "data": { + "id": "70045", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49417963", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "B-546C755A", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T15:40:49-04:00", + "departure_time": "2021-09-27T15:41:45-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 70 + }, + "id": "prediction-49417964-70045-70", + "relationships": { + "route": { + "data": { + "id": "Blue", + "type": "route" + } + }, + "stop": { + "data": { + "id": "70045", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49417964", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "B-546C7514", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T15:45:49-04:00", + "departure_time": "2021-09-27T15:46:45-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 70 + }, + "id": "prediction-49417965-70045-70", + "relationships": { + "route": { + "data": { + "id": "Blue", + "type": "route" + } + }, + "stop": { + "data": { + "id": "70045", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49417965", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "B-546C7357", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T15:50:49-04:00", + "departure_time": "2021-09-27T15:51:45-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 70 + }, + "id": "prediction-49417966-70045-70", + "relationships": { + "route": { + "data": { + "id": "Blue", + "type": "route" + } + }, + "stop": { + "data": { + "id": "70045", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49417966", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "B-546C7568", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T15:55:49-04:00", + "departure_time": "2021-09-27T15:56:45-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 70 + }, + "id": "prediction-49417967-70045-70", + "relationships": { + "route": { + "data": { + "id": "Blue", + "type": "route" + } + }, + "stop": { + "data": { + "id": "70045", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49417967", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "B-546C7336", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T16:00:49-04:00", + "departure_time": "2021-09-27T16:01:45-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 70 + }, + "id": "prediction-49417968-70045-70", + "relationships": { + "route": { + "data": { + "id": "Blue", + "type": "route" + } + }, + "stop": { + "data": { + "id": "70045", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49417968", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "B-546C756F", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T16:05:49-04:00", + "departure_time": "2021-09-27T16:06:45-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 70 + }, + "id": "prediction-49417969-70045-70", + "relationships": { + "route": { + "data": { + "id": "Blue", + "type": "route" + } + }, + "stop": { + "data": { + "id": "70045", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49417969", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "B-546C7358", + "type": "vehicle" + } + } + }, + "type": "prediction" + }, + { + "attributes": { + "arrival_time": "2021-09-27T16:10:49-04:00", + "departure_time": "2021-09-27T16:11:45-04:00", + "direction_id": 0, + "schedule_relationship": null, + "status": null, + "stop_sequence": 70 + }, + "id": "prediction-49417970-70045-70", + "relationships": { + "route": { + "data": { + "id": "Blue", + "type": "route" + } + }, + "stop": { + "data": { + "id": "70045", + "type": "stop" + } + }, + "trip": { + "data": { + "id": "49417970", + "type": "trip" + } + }, + "vehicle": { + "data": { + "id": "B-546C7556", + "type": "vehicle" + } + } + }, + "type": "prediction" + } + ], + "jsonapi": { + "version": "1.0" + } +} \ No newline at end of file diff --git a/test/test_mbta.py b/test/test_mbta.py new file mode 100644 index 0000000..a6f4210 --- /dev/null +++ b/test/test_mbta.py @@ -0,0 +1,54 @@ +import pytest +import pendulum + +from gtfs_realtime_translators.translators import MbtaGtfsRealtimeTranslator +from gtfs_realtime_translators.bindings import intersection_pb2 as intersection_gtfs_realtime +from gtfs_realtime_translators.factories import FeedMessage + + +@pytest.fixture +def mbta_subway(): + with open('test/fixtures/mbta_subway.json') as f: + raw = f.read() + return raw + +@pytest.fixture +def mbta_bus(): + with open('test/fixtures/mbta_bus.json') as f: + raw = f.read() + return raw + +def test_mbta_subway_realtime_arrival(mbta_subway): + translator = MbtaGtfsRealtimeTranslator() + message = translator(mbta_subway) + + entity = message.entity[0] + trip_update = entity.trip_update + stop_time_update = trip_update.stop_time_update[0] + + assert message.header.gtfs_realtime_version == FeedMessage.VERSION + + assert entity.id == '1' + assert entity.trip_update.trip.route_id == 'Blue' + assert entity.trip_update.trip.trip_id == '49417961' + assert stop_time_update.stop_id == '70045' + assert stop_time_update.arrival.time == 1632770790 + assert stop_time_update.departure.time == 1632770846 + +def test_mbta_bus_realtime_arrival(mbta_bus): + translator = MbtaGtfsRealtimeTranslator() + message = translator(mbta_bus) + + entity = message.entity[0] + trip_update = entity.trip_update + stop_time_update = trip_update.stop_time_update[0] + + assert message.header.gtfs_realtime_version == FeedMessage.VERSION + + assert entity.id == '3' + assert entity.trip_update.trip.route_id == '66' + assert entity.trip_update.trip.trip_id == '49181349' + assert stop_time_update.stop_id == '1357' + assert stop_time_update.arrival.time == 1632778733 + assert stop_time_update.departure.time == 1632778733 +