Remove stalled (cancelled) trains from extract_stop_dict (#30)

* Remove stalled (cancelled) trains from extract_stop_dict

* Run python3.6 CI on older Ubuntu

* Make stalled timeout configurable

* Support stalled-timeout flag in 'underground stops'

.github/workflows/push.yml

@@ -4,11 +4,15 @@ on: [push, pull_request]
 
 jobs:
   build:
-    runs-on: ubuntu-latest
+    runs-on: ${{matrix.os}}
     strategy:
       matrix:
-        python-version: ["3.6", "3.7", "3.8", "3.9", "3.10"]
-
+        os: [ubuntu-latest]
+        python-version: ["3.7", "3.8", "3.9", "3.10"]
+        include:
+          - python-version: "3.6"
+            os: "ubuntu-20.04"
+
     steps:
       - uses: actions/checkout@v3
 

src/underground/cli/stops.py

@@ -45,11 +45,20 @@ def datetime_to_epoch(dttm: datetime.datetime) -> int:
     default=metadata.DEFAULT_TIMEZONE,
     help="Output timezone. Ignored if --epoch. Default to NYC time.",
 )
-def main(route, fmt, retries, api_key, timezone):
+@click.option(
+    "-s",
+    "--stalled-timeout",
+    "stalled_timeout",
+    default=90,
+    help="Number of seconds between the last movement of a train and the API"
+    " update before considering a train stalled. Default is 90 as recommended"
+    " by the MTA. Numbers less than 1 disable this check.",
+)
+def main(route, fmt, retries, api_key, timezone, stalled_timeout):
     """Print out train departure times for all stops on a subway line."""
     stops = (
         SubwayFeed.get(api_key=api_key, route_or_url=route, retries=retries)
-        .extract_stop_dict(timezone=timezone)
+        .extract_stop_dict(timezone=timezone, stalled_timeout=stalled_timeout)
         .get(route, dict())
     )
 

src/underground/models.py

@@ -2,7 +2,6 @@
 
 import datetime
 import typing
-from operator import attrgetter
 
 import pydantic
 import pytz
@@ -212,30 +211,51 @@ def get(route_or_url: str, retries: int = 100, api_key: str = None) -> "SubwayFe
             )
         )
 
-    def extract_stop_dict(self, timezone: str = metadata.DEFAULT_TIMEZONE) -> dict:
+    def extract_stop_dict(
+        self, timezone: str = metadata.DEFAULT_TIMEZONE, stalled_timeout: int = 90
+    ) -> dict:
         """Get the departure times for all stops in the feed.
 
         Parameters
         ----------
         timezone : str
             Name of the timezone to return within. Default to NYC time.
+        stalled_timeout : int
+            Number of seconds between the last movement of a train and the API update before
+            considering a train stalled. Default is 90 as recommended by the MTA.
+            Numbers less than 1 disable this check.
 
         Returns
         -------
         dict
             Dictionary containing train departure for all stops in the gtfs data.
             The dictionary will be a schema like ``{route: {stop: [t1, t2]}}``.
-        
+
         """
-        # filter down data to trips with an update
-        entities_with_updates = filter(lambda x: x.trip_update is not None, self.entity)
-        trip_updates = map(attrgetter("trip_update"), entities_with_updates)
+
+        trip_updates = (x.trip_update for x in self.entity if x.trip_update is not None)
+        vehicles = {
+            e.vehicle.trip.trip_id: e.vehicle
+            for e in self.entity
+            if e.vehicle is not None
+        }
+
+        def is_trip_active(update: TripUpdate) -> bool:
+            has_route = update.trip.route_is_assigned
+            has_stops = update.stop_time_update is not None
+
+            vehicle = vehicles.get(update.trip.trip_id)
+            if stalled_timeout < 1 or vehicle is None or vehicle.timestamp is None:
+                return has_route and has_stops
+
+            # as recommended by the MTA, we use these timestamps to determine if a train is stalled
+            train_stalled = (
+                self.header.timestamp - vehicle.timestamp
+            ) > datetime.timedelta(seconds=stalled_timeout)
+            return has_route and has_stops and not train_stalled
 
         # grab the updates with routes and stop times
-        trip_updates_with_stops = filter(
-            lambda x: x.trip.route_is_assigned and x.stop_time_update is not None,
-            trip_updates,
-        )
+        trip_updates_with_stops = filter(is_trip_active, trip_updates)
         # create (route, stop, time) tuples from each trip
         stops_flat = (
             (

src/underground/version

@@ -1 +1 @@
-0.3.3
+0.3.4

test/test_models.py

@@ -147,6 +147,91 @@ def test_extract_dict_elapsed_ignored():
     assert "EXISTS" in stops["1"]
 
 
+def test_extract_dict_stalled_train_omitted():
+    """Test that a stalled train is omitted from the stop extraction."""
+    sample_data = {
+        "header": {"gtfs_realtime_version": "1.0", "timestamp": 1699239196},
+        "entity": [
+            {
+                "id": "000021F",
+                "trip_update": {
+                    "trip": {
+                        "trip_id": "128000_F..S",
+                        "start_time": "21:20:00",
+                        "start_date": "20231105",
+                        "route_id": "F",
+                    },
+                    "stop_time_update": [
+                        {
+                            "arrival": {"time": 1699239913},
+                            "departure": {"time": 1699239913},
+                            "stop_id": "D17S",
+                        }
+                    ],
+                },
+            },
+            {
+                "id": "000022F",
+                "vehicle": {
+                    "trip": {
+                        "trip_id": "128000_F..S",
+                        "start_time": "21:20:00",
+                        "start_date": "20231105",
+                        "route_id": "F",
+                    },
+                    "current_stop_sequence": 13,
+                    "current_status": 1,
+                    "timestamp": 1699238728,  # STALLED TRAIN
+                    "stop_id": "G14S",
+                },
+            },
+            {
+                "id": "000025F",
+                "trip_update": {
+                    "trip": {
+                        "trip_id": "129300_F..S",
+                        "start_time": "21:33:00",
+                        "start_date": "20231105",
+                        "route_id": "F",
+                    },
+                    "stop_time_update": [
+                        {
+                            "arrival": {"time": 1699240511},
+                            "departure": {"time": 1699240511},
+                            "stop_id": "D17S",
+                        }
+                    ],
+                },
+            },
+            {
+                "id": "000026F",
+                "vehicle": {
+                    "trip": {
+                        "trip_id": "129300_F..S",
+                        "start_time": "21:33:00",
+                        "start_date": "20231105",
+                        "route_id": "F",
+                    },
+                    "current_stop_sequence": 10,
+                    "current_status": 1,
+                    "timestamp": 1699239183,
+                    "stop_id": "G11S",
+                },
+            },
+        ],
+    }
+    feed = SubwayFeed(**sample_data)
+
+    stops_no_timeout = feed.extract_stop_dict(stalled_timeout=0)
+    assert len(stops_no_timeout["F"]["D17S"]) == 2
+
+    stops_default = feed.extract_stop_dict()  # stalled_timeout=90
+    assert len(stops_default["F"]["D17S"]) == 1
+
+    stops_long_timeout = feed.extract_stop_dict(stalled_timeout=900)
+    assert len(stops_long_timeout["F"]["D17S"]) == 2
+
+
 def test_empty_route_id():
     """Test the route functionality when the route id is a blacnk string.