Introduce CTA Subway and Bus Translators (#15)

Intersection · Oct 10, 2019 · 1cd906a · 1cd906a
1 parent babdd7c
commit 1cd906a
Show file tree

Hide file tree

Showing 8 changed files with 293 additions and 1 deletion.
diff --git a/gtfs_realtime_translators/registry/registry.py b/gtfs_realtime_translators/registry/registry.py
@@ -1,7 +1,8 @@
 import warnings
 
 from gtfs_realtime_translators.translators import LaMetroGtfsRealtimeTranslator, \
-        SeptaRegionalRailTranslator, MtaSubwayGtfsRealtimeTranslator, NjtRailGtfsRealtimeTranslator
+    SeptaRegionalRailTranslator, MtaSubwayGtfsRealtimeTranslator, NjtRailGtfsRealtimeTranslator, \
+    CtaSubwayGtfsRealtimeTranslator, CtaBusGtfsRealtimeTranslator
 
 
 class TranslatorKeyWarning(Warning):
@@ -12,6 +13,8 @@ class TranslatorRegistry:
     TRANSLATORS = {
         'la-metro': LaMetroGtfsRealtimeTranslator,
         'septa-regional-rail': SeptaRegionalRailTranslator,
+        'cta-subway': CtaSubwayGtfsRealtimeTranslator,
+        'cta-bus': CtaBusGtfsRealtimeTranslator,
         'mta-subway': MtaSubwayGtfsRealtimeTranslator,
         'njt-rail': NjtRailGtfsRealtimeTranslator,
     }

diff --git a/gtfs_realtime_translators/translators/__init__.py b/gtfs_realtime_translators/translators/__init__.py
@@ -2,3 +2,5 @@
 from .septa_regional_rail import SeptaRegionalRailTranslator
 from .mta_subway import MtaSubwayGtfsRealtimeTranslator
 from .njt_rail import NjtRailGtfsRealtimeTranslator
+from .cta_subway import CtaSubwayGtfsRealtimeTranslator
+from .cta_bus import CtaBusGtfsRealtimeTranslator
diff --git a/gtfs_realtime_translators/translators/cta_bus.py b/gtfs_realtime_translators/translators/cta_bus.py
@@ -0,0 +1,45 @@
+import json
+
+import pendulum
+
+from gtfs_realtime_translators.factories import TripUpdate, FeedMessage
+
+
+class CtaBusGtfsRealtimeTranslator:
+    TIMEZONE = 'America/Chicago'
+
+    def __call__(self, data):
+        json_data = json.loads(data)
+        predictions = json_data['bustime-response']['prd']
+        entities = [self.__make_trip_update(idx, arr) for idx, arr in enumerate(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_update(cls, _id, prediction):
+        entity_id = str(_id + 1)
+        route_id = prediction['rt']
+        stop_id = prediction['stpid']
+
+        # Per the docs, the prediction is either for a scheduled departure, or realtime arrival
+        departure_type = 'D'
+        is_scheduled = prediction['typ'] == departure_type
+
+        parsed_arrival_time = cls.__to_unix_time(prediction['prdtm'])
+        arrival_time = None if is_scheduled else parsed_arrival_time
+
+        ##### Intersection Extensions
+        headsign = prediction['des']
+        scheduled_arrival_time = parsed_arrival_time if is_scheduled else None
+
+        return TripUpdate.create(entity_id=entity_id,
+                                route_id=route_id,
+                                stop_id=stop_id,
+                                arrival_time=arrival_time,
+                                headsign=headsign,
+                                scheduled_arrival_time=scheduled_arrival_time)
diff --git a/gtfs_realtime_translators/translators/cta_subway.py b/gtfs_realtime_translators/translators/cta_subway.py
@@ -0,0 +1,42 @@
+import json
+
+import pendulum
+
+from gtfs_realtime_translators.factories import TripUpdate, FeedMessage
+
+
+class CtaSubwayGtfsRealtimeTranslator:
+    TIMEZONE = 'America/Chicago'
+
+    def __call__(self, data):
+        json_data = json.loads(data)
+        prediction = json_data['ctatt']['eta']
+        entities = [self.__make_trip_update(idx, arr) for idx, arr in enumerate(prediction)]
+
+        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_update(cls, _id, prediction):
+        entity_id = str(_id + 1)
+        route_id = prediction['rt']
+        stop_id = prediction['stpId']
+
+        is_scheduled = prediction['isSch'] == '1'
+        parsed_arrival_time = cls.__to_unix_time(prediction['arrT'])
+        arrival_time = None if is_scheduled else parsed_arrival_time
+
+        ##### Intersection Extensions
+        headsign = prediction['destNm']
+        scheduled_arrival_time = parsed_arrival_time if is_scheduled else None
+
+        return TripUpdate.create(entity_id=entity_id,
+                                route_id=route_id,
+                                stop_id=stop_id,
+                                arrival_time=arrival_time,
+                                headsign=headsign,
+                                scheduled_arrival_time=scheduled_arrival_time)
diff --git a/test/fixtures/cta_bus.json b/test/fixtures/cta_bus.json
@@ -0,0 +1,42 @@
+{
+	"bustime-response": {
+		"prd": [
+			{
+				"tmstmp": "20191008 10:34",
+				"typ": "A",
+				"stpnm": "Sheridan & Arthur (Red Line)",
+				"stpid": "1203",
+				"vid": "4379",
+				"dstp": 2603,
+				"rt": "147",
+				"rtdd": "147",
+				"rtdir": "Northbound",
+				"des": "Howard Station",
+				"prdtm": "20191008 10:38",
+				"tablockid": "147 -501",
+				"tatripid": "1063847",
+				"dly": false,
+				"prdctdn": "3",
+				"zone": ""
+			},
+			{
+				"tmstmp": "20191008 10:26",
+				"typ": "D",
+				"stpnm": "Sheridan & Arthur (Red Line)",
+				"stpid": "1203",
+				"vid": "1396",
+				"dstp": 13893,
+				"rt": "155",
+				"rtdd": "155",
+				"rtdir": "Eastbound",
+				"des": "Morse Red Line",
+				"prdtm": "20191008 10:45",
+				"tablockid": "155 -502",
+				"tatripid": "1005547",
+				"dly": false,
+				"prdctdn": "10",
+				"zone": ""
+			}
+		]
+	}
+}
diff --git a/test/fixtures/cta_subway.json b/test/fixtures/cta_subway.json
@@ -0,0 +1,51 @@
+{
+    "ctatt": {
+        "tmst": "2019-10-07T14:29:33",
+        "errCd": "0",
+        "errNm": null,
+        "eta": [
+            {
+                "staId": "41300",
+                "stpId": "30251",
+                "staNm": "Loyola",
+                "stpDe": "Service toward Howard",
+                "rn": "806",
+                "rt": "Red",
+                "destSt": "30173",
+                "destNm": "Howard",
+                "trDr": "1",
+                "prdt": "2019-10-07T14:29:02",
+                "arrT": "2019-10-07T14:30:02",
+                "isApp": "1",
+                "isSch": "0",
+                "isDly": "0",
+                "isFlt": "0",
+                "flags": null,
+                "lat": "41.99673",
+                "lon": "-87.65923",
+                "heading": "357"
+            },
+            {
+                "staId": "41300",
+                "stpId": "30251",
+                "staNm": "Loyola",
+                "stpDe": "Service toward Howard",
+                "rn": "824",
+                "rt": "Red",
+                "destSt": "30173",
+                "destNm": "Howard",
+                "trDr": "1",
+                "prdt": "2019-10-07T14:28:56",
+                "arrT": "2019-10-07T14:32:56",
+                "isApp": "0",
+                "isSch": "1",
+                "isDly": "0",
+                "isFlt": "0",
+                "flags": null,
+                "lat": "41.9835",
+                "lon": "-87.65884",
+                "heading": "358"
+            }
+        ]
+    }
+}
diff --git a/test/test_cta_bus.py b/test/test_cta_bus.py
@@ -0,0 +1,53 @@
+import pytest
+import pendulum
+
+from gtfs_realtime_translators.translators import CtaBusGtfsRealtimeTranslator
+from gtfs_realtime_translators.bindings import intersection_pb2 as intersection_gtfs_realtime
+from gtfs_realtime_translators.factories import FeedMessage
+
+
+@pytest.fixture
+def cta_bus():
+    with open('test/fixtures/cta_bus.json') as f:
+        raw = f.read()
+    return raw
+
+
+def test_cta_bus_realtime_arrival(cta_bus):
+    translator = CtaBusGtfsRealtimeTranslator()
+    with pendulum.test(pendulum.datetime(2019, 2, 20, 17)):
+        message = translator(cta_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 == '1'
+    assert entity.trip_update.trip.route_id == '147'
+    assert stop_time_update.stop_id == '1203'
+    assert stop_time_update.arrival.time == 1570531080
+
+    # Test Intersection extensions
+    intersection_trip_update = trip_update.Extensions[intersection_gtfs_realtime.intersection_trip_update]
+    assert intersection_trip_update.headsign == 'Howard Station'
+
+    intersection_stop_time_update = stop_time_update.Extensions[intersection_gtfs_realtime.intersection_stop_time_update]
+    assert intersection_stop_time_update.scheduled_arrival.time == 0
+
+def test_cta_bus_scheduled_departure(cta_bus):
+    translator = CtaBusGtfsRealtimeTranslator()
+    with pendulum.test(pendulum.datetime(2019, 2, 20, 17)):
+        message = translator(cta_bus)
+
+    entity = message.entity[1]
+    trip_update = entity.trip_update
+    stop_time_update = trip_update.stop_time_update[0]
+
+    assert entity.id == '2'
+    assert stop_time_update.arrival.time == 0
+
+    # Test Intersection extensions
+    intersection_stop_time_update = stop_time_update.Extensions[intersection_gtfs_realtime.intersection_stop_time_update]
+    assert intersection_stop_time_update.scheduled_arrival.time == 1570531500
diff --git a/test/test_cta_subway.py b/test/test_cta_subway.py
@@ -0,0 +1,54 @@
+import pytest
+import pendulum
+
+from gtfs_realtime_translators.translators import CtaSubwayGtfsRealtimeTranslator
+from gtfs_realtime_translators.bindings import intersection_pb2 as intersection_gtfs_realtime
+from gtfs_realtime_translators.factories import FeedMessage
+
+
+@pytest.fixture
+def cta_subway():
+    with open('test/fixtures/cta_subway.json') as f:
+        raw = f.read()
+
+    return raw
+
+
+def test_cta_subway_realtime_arrival(cta_subway):
+    translator = CtaSubwayGtfsRealtimeTranslator()
+    with pendulum.test(pendulum.datetime(2019, 2, 20, 17)):
+        message = translator(cta_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 == 'Red'
+    assert stop_time_update.stop_id == '30251'
+    assert stop_time_update.arrival.time == 1570458602
+
+    # Test Intersection extensions
+    intersection_trip_update = trip_update.Extensions[intersection_gtfs_realtime.intersection_trip_update]
+    assert intersection_trip_update.headsign == 'Howard'
+
+    intersection_stop_time_update = stop_time_update.Extensions[intersection_gtfs_realtime.intersection_stop_time_update]
+    assert intersection_stop_time_update.scheduled_arrival.time == 0
+
+def test_cta_subway_scheduled_arrival(cta_subway):
+    translator = CtaSubwayGtfsRealtimeTranslator()
+    with pendulum.test(pendulum.datetime(2019, 2, 20, 17)):
+        message = translator(cta_subway)
+
+    entity = message.entity[1]
+    trip_update = entity.trip_update
+    stop_time_update = trip_update.stop_time_update[0]
+
+    assert entity.id == '2'
+    assert stop_time_update.arrival.time == 0
+
+    # Test Intersection extensions
+    intersection_stop_time_update = stop_time_update.Extensions[intersection_gtfs_realtime.intersection_stop_time_update]
+    assert intersection_stop_time_update.scheduled_arrival.time == 1570458776