TRAN-7701: add custom_status to cta_subway; and scheduled_interval fi…

…eld to proto (#61)

* TRAN-7701: add custom_status to cta_subway

* TRAN-7701: add scheduled interval field to intersection extensions

* TRAN-7701: filter stops in feed

gtfs_realtime_translators/bindings/intersection.proto

@@ -11,6 +11,7 @@ message IntersectionTripUpdate {
   optional string block_id = 6;
   optional string agency_timezone = 7;
   optional string custom_status = 8;
+  optional int32 scheduled_interval = 9;
 }
 
 message IntersectionStopTimeUpdate {

gtfs_realtime_translators/factories/factories.py

@@ -10,7 +10,7 @@ def create(entity_id, **kwargs):
 
 
 class TripUpdate:
-    
+
     @staticmethod
     def __set_stop_time_events(arrival_time, departure_time):
         arrival = gtfs_realtime.TripUpdate.StopTimeEvent(time=arrival_time)
@@ -60,6 +60,7 @@ def create(*args, **kwargs):
         block_id = kwargs.get('block_id', None)
         agency_timezone = kwargs.get('agency_timezone', None)
         custom_status = kwargs.get('custom_status', None)
+        scheduled_interval = kwargs.get('scheduled_interval', None)
 
         trip_descriptor = gtfs_realtime.TripDescriptor(trip_id=trip_id,
                                                        route_id=route_id,
@@ -97,6 +98,8 @@ def create(*args, **kwargs):
             trip_update.Extensions[intersection_gtfs_realtime.intersection_trip_update].agency_timezone = agency_timezone
         if custom_status:
             trip_update.Extensions[intersection_gtfs_realtime.intersection_trip_update].custom_status = custom_status
+        if scheduled_interval:
+            trip_update.Extensions[intersection_gtfs_realtime.intersection_trip_update].scheduled_interval = scheduled_interval
 
         return Entity.create(entity_id,
                              trip_update=trip_update)

gtfs_realtime_translators/translators/cta_subway.py

@@ -8,10 +8,22 @@
 class CtaSubwayGtfsRealtimeTranslator:
     TIMEZONE = 'America/Chicago'
 
+    def __init__(self, **kwargs):
+        stop_list = kwargs.get('stop_list')
+        if stop_list:
+            self.stop_list = stop_list.split(',')
+        else:
+            self.stop_list = None
+
     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)]
+        predictions = json_data['ctatt']['eta']
+
+        entities = []
+        for idx, prediction in enumerate(predictions):
+            stop_id = prediction['stpId']
+            if not self.stop_list or stop_id in self.stop_list:
+                entities.append(self.__make_trip_update(idx, prediction))
 
         return FeedMessage.create(entities=entities)
 
@@ -34,9 +46,32 @@ def __make_trip_update(cls, _id, prediction):
         headsign = prediction['destNm']
         scheduled_arrival_time = parsed_arrival_time if is_scheduled else None
 
+        parsed_prediction_time = cls.__to_unix_time(prediction['prdt'])
+        custom_status = cls.__get_custom_status(parsed_arrival_time,
+                                                parsed_prediction_time)
+        scheduled_interval = cls.__get_scheduled_interval(is_scheduled,
+                                                          prediction['schInt'])
+
         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)
+                                 route_id=route_id,
+                                 stop_id=stop_id,
+                                 arrival_time=arrival_time,
+                                 headsign=headsign,
+                                 scheduled_arrival_time=scheduled_arrival_time,
+                                 custom_status=custom_status,
+                                 agency_timezone=cls.TIMEZONE,
+                                 scheduled_interval=scheduled_interval)
+
+    @classmethod
+    def __get_custom_status(cls, arrival_time, prediction_time):
+        minutes_until_train_arrives = (arrival_time - prediction_time) / 60
+        if minutes_until_train_arrives <= 1:
+            return 'DUE'
+        return f'{round(minutes_until_train_arrives)} min'
+
+    @classmethod
+    def __get_scheduled_interval(cls, is_scheduled, scheduled_interval):
+        if is_scheduled:
+            scheduled_interval_seconds = int(scheduled_interval) * 60
+            return scheduled_interval_seconds
+        return None