Feat: custom status field and NJT/Amtrak route updates (#14)

* Add custom status field

* Fix routes

* Update route logic for amtrak

* Amtrak short name fix

* Add prefix change and test

* Refactoring route logic and timezone

* Handle missing origin and destination

* Refactor route_id logic

* Use class level constant for timezone

* Minor refactor and typo fix

gtfs_realtime_translators/bindings/intersection.proto

@@ -10,6 +10,7 @@ message IntersectionTripUpdate {
   optional string route_text_color = 5;
   optional string block_id = 6;
   optional string agency_timezone = 7;
+  optional string custom_status = 8;
 }
 
 message IntersectionStopTimeUpdate {

gtfs_realtime_translators/factories/factories.py

@@ -42,6 +42,7 @@ def create(*args, **kwargs):
         route_text_color = kwargs.get('route_text_color', None)
         block_id = kwargs.get('block_id', None)
         agency_timezone = kwargs.get('agency_timezone', None)
+        custom_status = kwargs.get('custom_status', None)
 
         trip_descriptor = gtfs_realtime.TripDescriptor(trip_id=trip_id,
                                                        route_id=route_id)
@@ -76,6 +77,8 @@ def create(*args, **kwargs):
             trip_update.Extensions[intersection_gtfs_realtime.intersection_trip_update].block_id = block_id
         if agency_timezone:
             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
 
         return Entity.create(entity_id,
                              trip_update=trip_update)

gtfs_realtime_translators/translators/njt_rail.py

@@ -18,14 +18,17 @@ class NjtRailGtfsRealtimeTranslator:
 
     https://usermanual.wiki/Document/NJTRANSIT20REAL20Time20Data20Interface20Instructions2020Ver2025.785373145.pdf
     """
+
+    TIMEZONE = 'America/New_York'
+
     def __call__(self, data):
         station_data = xmltodict.parse(data)
         entities = self.__make_trip_updates(station_data)
         return FeedMessage.create(entities=entities)
 
     @classmethod
     def __to_unix_time(cls, time):
-        datetime = pendulum.from_format(time, 'DD-MMM-YYYY HH:mm:ss A', tz='America/New_York').in_tz('UTC')
+        datetime = pendulum.from_format(time, 'DD-MMM-YYYY HH:mm:ss A', tz=cls.TIMEZONE).in_tz('UTC')
         return datetime
 
     @classmethod
@@ -35,11 +38,10 @@ def __make_trip_updates(cls, data):
         station_data_item = data['STATION']['ITEMS'].values()
         for value in station_data_item:
             for idx, item_entry in enumerate(value):
-                route_id = None
 
                 # Intersection Extensions
                 headsign = item_entry['DESTINATION']
-                route_short_name = item_entry['LINEABBREVIATION']
+                route_short_name = cls.__get_route_short_name(item_entry)
                 route_long_name = cls.__get_route_long_name(item_entry)
                 route_color = item_entry['BACKCOLOR']
                 route_text_color = item_entry['FORECOLOR']
@@ -50,14 +52,13 @@ def __make_trip_updates(cls, data):
                 scheduled_datetime = cls.__to_unix_time(item_entry['SCHED_DEP_DATE'])
                 departure_time = int(scheduled_datetime.add(seconds=int(item_entry['SEC_LATE'])).timestamp())
                 scheduled_departure_time = int(scheduled_datetime.timestamp())
+                custom_status = item_entry['STATUS']
 
+                origin_and_destination = None
                 for stop in item_entry['STOPS'].values():
                     origin_and_destination = [stop[i] for i in (0, -1)]
 
-                    route_id = cls.__get_route_id(line=item_entry['LINE'],
-                                                  line_abbreviation=item_entry['LINEABBREVIATION'],
-                                                  origin=origin_and_destination[0],
-                                                  destination=origin_and_destination[1])
+                route_id = cls.__get_route_id(item_entry, origin_and_destination)
 
                 trip_update = TripUpdate.create(entity_id=str(idx + 1),
                                                 departure_time=departure_time,
@@ -74,65 +75,94 @@ def __make_trip_updates(cls, data):
                                                 headsign=headsign,
                                                 track=track,
                                                 block_id=block_id,
-                                                agency_timezone='America/New_York')
+                                                agency_timezone=cls.TIMEZONE,
+                                                custom_status=custom_status)
                 trip_updates.append(trip_update)
 
         return trip_updates
 
     @classmethod
-    def __get_route_id(cls, **data):
+    def __get_route_id(cls, data, origin_and_destination):
         """
-        This function resolves route_ids for NJT. The logic of determining a route_id based
-        on origin or destination is necessary to discern multiple routes that are mapped to the same line.
+        This function resolves route_ids for NJT.
 
-        For instance, the North Jersey Coastline line operates two different routes. All trains with an
-        origin or destination of New York Penn Station should resolve to route_id 10 and the others route_id 11
+        The algorithm is as follows:
+        1) Try to get the route_id based on the line name (or line abbreviation for Amtrak), otherwise...
+        2) Try to get the route_id based on the origin and destination, otherwise return None
+
+        For #2, this logic is necessary to discern multiple routes that are mapped to the same line. For instance, the
+        North Jersey Coast Line operates two different routes. All trains with an origin or destination
+        of New York Penn Station should resolve to route_id 10 and the others route_id 11
 
         :param data: keyword args containing data needed to perform the route logic
+        :param origin_and_destination: an array containing the origin at index 0 and destination at index 1
         :return: route_id
         """
+
+        route_id = cls.__get_route_id_by_line_data(data)
+        if route_id:
+            return route_id
+        if origin_and_destination:
+            return cls.__get_route_id_by_origin_or_destination(data, origin_and_destination)
+        return None
+
+    @classmethod
+    def __get_route_id_by_origin_or_destination(cls, data, origin_and_destination):
+        origin = origin_and_destination[0]
+        destination = origin_and_destination[1]
+        origin_name = origin['NAME'].replace(' ', '_').lower()
+        destination_name = destination['NAME'].replace(' ', '_').lower()
+
+        key = data['LINE'].replace(' ', '_').lower()
+        if key == 'montclair-boonton_line':
+            hoboken = 'hoboken'
+            origins_and_destinations = {'denville', 'dover', 'mount_olive', 'lake_hopatcong', 'hackettstown'}
+            if origin_name == hoboken and destination_name in origins_and_destinations:
+                return '2'
+            if origin_name in origins_and_destinations and destination_name == hoboken:
+                return '2'
+            return '3'
+
+        if key == 'north_jersey_coast_line':
+            origins_and_destinations = {'new_york_penn_station'}
+            if origin_name in origins_and_destinations or destination_name in origins_and_destinations:
+                return '10'
+            return '11'
+        return None
+
+    @classmethod
+    def __get_route_id_by_line_data(cls, data):
         route_id_lookup = {
             'atlantic_city_line': '1',
-            'montclair-boonton_line': None,
             'main_line': '5',
             'bergen_county_line': '6',
             'morristown_line': '7',
             'gladstone_branch': '8',
             'northeast_corridor_line': '9',
-            'north_jersey_coast_line': None,
+            'pascack_valley_line': '13',
+            'princeton_shuttle': '14',
+            'raritan_valley_line': '15',
+            'meadowlands_rail_line': '17',
         }
 
-        key = 'amtrak' if data['line_abbreviation'] == 'AMTK' else data['line'].replace(' ', '_').lower()
-        route_id = route_id_lookup.get(key, None)
-        if route_id is not None:
-            return route_id
+        amtrak_route_id = 'AMTK'
+        if data['LINEABBREVIATION'] == amtrak_route_id:
+            return amtrak_route_id
 
-        def get_route_id_by_origin_or_destination(line_key, line_name, origin, destination):
-            origin_name = origin['NAME'].replace(' ', '_').lower()
-            destination_name = destination['NAME'].replace(' ', '_').lower()
-
-            if line_key == 'montclair-boonton_line':
-                hoboken = 'hoboken'
-                origins_and_destinations = {'denville', 'dover', 'mount_olive', 'lake_hopatcong', 'hackettstown'}
-                if origin_name == hoboken and destination_name in origins_and_destinations:
-                    return '2'
-                if origin_name in origins_and_destinations and destination_name == hoboken:
-                    return '2'
-                return '3'
-
-            if line_key == 'north_jersey_coast_line':
-                origins_and_destinations = {'new_york_penn_station'}
-                if origin_name in origins_and_destinations or destination_name in origins_and_destinations:
-                    return '10'
-                return '11'
-            if line_key == 'amtrak':
-                return line_name
-            return None
-
-        return get_route_id_by_origin_or_destination(key, data['line'], data['origin'], data['destination'])
+        key = data['LINE'].replace(' ', '_').lower()
+        route_id = route_id_lookup.get(key, None)
+        return route_id if route_id else None
 
     @classmethod
     def __get_route_long_name(cls, data):
-        if data['LINEABBREVIATION'] == 'AMTK':
-            return f"Amtrak {data['LINE']}".title()
+        amtrak_prefix = 'AMTRAK'
+        abbreviation = data['LINEABBREVIATION']
+        if abbreviation == 'AMTK':
+            return amtrak_prefix.title() if data['LINE'] == amtrak_prefix else f"Amtrak {data['LINE']}".title()
         return data['LINE']
+
+    @classmethod
+    def __get_route_short_name(cls, data):
+        if data['LINEABBREVIATION'] == 'AMTK':
+            return data['LINE']
+        return data['LINEABBREVIATION']

test/fixtures/njt_rail.xml

@@ -9,7 +9,7 @@
             <SCHED_DEP_DATE>02-Oct-2019 03:02:00 PM</SCHED_DEP_DATE>
             <DESTINATION>Boston</DESTINATION>
             <TRACK>2</TRACK>
-            <LINE>REGIONAL</LINE>
+            <LINE>AMTRAK</LINE>
             <TRAIN_ID>A176</TRAIN_ID>
             <CONNECTING_TRAIN_ID>
             </CONNECTING_TRAIN_ID>
@@ -659,7 +659,7 @@
             <TRAIN_ID>3154</TRAIN_ID>
             <CONNECTING_TRAIN_ID>
             </CONNECTING_TRAIN_ID>
-            <STATUS>CANCELLED</STATUS>
+            <STATUS> </STATUS>
             <SEC_LATE>0</SEC_LATE>
             <LAST_MODIFIED>02-Oct-2019 02:51:30 PM</LAST_MODIFIED>
             <BACKCOLOR>black</BACKCOLOR>

test/test_njt_rail.py

@@ -38,6 +38,7 @@ def test_njt_data(njt_rail):
     assert intersection_trip_update.route_text_color == 'white'
     assert intersection_trip_update.block_id == '3154'
     assert intersection_trip_update.agency_timezone == 'America/New_York'
+    assert intersection_trip_update.custom_status == ''
 
     intersection_stop_time_update = stop_time_update.Extensions[intersection_gtfs_realtime.intersection_stop_time_update]
     assert intersection_stop_time_update.track == '1'
@@ -58,22 +59,52 @@ def test_njt_data_amtrak(njt_rail):
     assert entity.id == '1'
 
     assert trip_update.trip.trip_id == ''
-    assert trip_update.trip.route_id == 'REGIONAL'
+    assert trip_update.trip.route_id == 'AMTK'
 
     assert stop_time_update.stop_id == 'NP'
     assert stop_time_update.departure.time == 1570044525
     assert stop_time_update.arrival.time == 1570044525
 
     intersection_trip_update = trip_update.Extensions[intersection_gtfs_realtime.intersection_trip_update]
     assert intersection_trip_update.headsign == 'Boston'
-    assert intersection_trip_update.route_short_name == 'AMTK'
-    assert intersection_trip_update.route_long_name == 'Amtrak Regional'
+    assert intersection_trip_update.route_short_name == 'AMTRAK'
+    assert intersection_trip_update.route_long_name == 'Amtrak'
     assert intersection_trip_update.route_color == 'yellow'
     assert intersection_trip_update.route_text_color == 'black'
     assert intersection_trip_update.block_id == 'A176'
     assert intersection_trip_update.agency_timezone == 'America/New_York'
+    assert intersection_trip_update.custom_status == 'All Aboard'
 
     intersection_stop_time_update = stop_time_update.Extensions[intersection_gtfs_realtime.intersection_stop_time_update]
     assert intersection_stop_time_update.track == '2'
     assert intersection_stop_time_update.scheduled_arrival.time == 1570042920
     assert intersection_stop_time_update.scheduled_departure.time == 1570042920
+
+    entity = message.entity[15]
+    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 == '16'
+
+    assert trip_update.trip.trip_id == ''
+    assert trip_update.trip.route_id == 'AMTK'
+
+    assert stop_time_update.stop_id == 'NP'
+    assert stop_time_update.departure.time == 1570047420
+    assert stop_time_update.arrival.time == 1570047420
+
+    intersection_trip_update = trip_update.Extensions[intersection_gtfs_realtime.intersection_trip_update]
+    assert intersection_trip_update.headsign == 'Washington'
+    assert intersection_trip_update.route_short_name == 'ACELA EXPRESS'
+    assert intersection_trip_update.route_long_name == 'Amtrak Acela Express'
+    assert intersection_trip_update.route_color == 'yellow'
+    assert intersection_trip_update.route_text_color == 'black'
+    assert intersection_trip_update.block_id == 'A2165'
+    assert intersection_trip_update.agency_timezone == 'America/New_York'
+    assert intersection_trip_update.custom_status == ''
+
+    intersection_stop_time_update = stop_time_update.Extensions[intersection_gtfs_realtime.intersection_stop_time_update]
+    assert intersection_stop_time_update.track == '3'
+    assert intersection_stop_time_update.scheduled_arrival.time == 1570047420
+    assert intersection_stop_time_update.scheduled_departure.time == 1570047420