Merge pull request #70 from Intersection/TRAN-8345

Intersection · Nov 25, 2024 · aa84493 · aa84493
2 parents ac3fb40 + 9530fb5
commit aa84493
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 3 deletions.
diff --git a/gtfs_realtime_translators/bindings/intersection.proto b/gtfs_realtime_translators/bindings/intersection.proto
@@ -29,3 +29,12 @@ extend transit_realtime.TripUpdate {
 extend transit_realtime.TripUpdate.StopTimeUpdate {
   optional IntersectionStopTimeUpdate intersection_stop_time_update = 1987;
 }
+
+
+message IntersectionVehicleDescriptor {
+  optional uint32 run_number = 1;
+}
+
+extend transit_realtime.VehicleDescriptor {
+  optional IntersectionVehicleDescriptor intersection_vehicle_descriptor = 1987;
+}
diff --git a/gtfs_realtime_translators/bindings/intersection_pb2.py b/gtfs_realtime_translators/bindings/intersection_pb2.py
diff --git a/gtfs_realtime_translators/factories/factories.py b/gtfs_realtime_translators/factories/factories.py
@@ -62,6 +62,7 @@ def create(*args, **kwargs):
         custom_status = kwargs.get('custom_status', None)
         scheduled_interval = kwargs.get('scheduled_interval', None)
         route_icon = kwargs.get('route_icon', None)
+        run_number = kwargs.get('run_number', None)
 
         trip_descriptor = gtfs_realtime.TripDescriptor(trip_id=trip_id,
                                                        route_id=route_id,
@@ -70,6 +71,8 @@ def create(*args, **kwargs):
         stop_time_update = gtfs_realtime.TripUpdate.StopTimeUpdate(arrival=arrival,
                                                                    departure=departure,
                                                                    stop_id=stop_id)
+
+        vehicle_descriptor = gtfs_realtime.VehicleDescriptor()
 
         if track:
             stop_time_update.Extensions[intersection_gtfs_realtime.intersection_stop_time_update].track = track
@@ -79,9 +82,12 @@ def create(*args, **kwargs):
             stop_time_update.Extensions[intersection_gtfs_realtime.intersection_stop_time_update].scheduled_departure.time = scheduled_departure
         if stop_name:
             stop_time_update.Extensions[intersection_gtfs_realtime.intersection_stop_time_update].stop_name = stop_name
+        if run_number:
+            vehicle_descriptor.Extensions[intersection_gtfs_realtime.intersection_vehicle_descriptor].run_number = run_number
 
         trip_update = gtfs_realtime.TripUpdate(trip=trip_descriptor,
-                                               stop_time_update=[stop_time_update])
+                                               stop_time_update=[stop_time_update],
+                                               vehicle=vehicle_descriptor)
 
         if headsign:
             trip_update.Extensions[intersection_gtfs_realtime.intersection_trip_update].headsign = headsign

diff --git a/gtfs_realtime_translators/translators/cta_subway.py b/gtfs_realtime_translators/translators/cta_subway.py
@@ -53,6 +53,7 @@ def __make_trip_update(cls, _id, prediction):
                                                           prediction['schInt'])
 
         route_icon = cls.__get_route_icon(prediction['flags'], headsign)
+        run_number = prediction['rn']
 
         return TripUpdate.create(entity_id=entity_id,
                                  route_id=route_id,
@@ -63,7 +64,8 @@ def __make_trip_update(cls, _id, prediction):
                                  custom_status=custom_status,
                                  agency_timezone=cls.TIMEZONE,
                                  scheduled_interval=scheduled_interval,
-                                 route_icon = route_icon)
+                                 route_icon=route_icon,
+                                 run_number=run_number)
 
     @classmethod
     def __get_custom_status(cls, arrival_time, prediction_time):