Add MBTA translator (#38)

Intersection · Sep 28, 2021 · 175c2ce · 175c2ce
1 parent 516467e
commit 175c2ce
Show file tree

Hide file tree

Showing 6 changed files with 801 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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
@@ -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"
+  }
+}