feat[pmtiles]: support multiple shapes per routes (#1365)

Author: davidgamez

functions-python/helpers/tests/test_transform.py

@@ -8,6 +8,7 @@
     to_float,
     get_safe_value,
     get_safe_float,
+    get_safe_int,
 )
 
 
@@ -154,3 +155,21 @@ def test_default_value(self):
         self.assertEqual(get_safe_float(row, "value", default_value=4.56), 4.56)
         row = {"value": None}
         self.assertEqual(get_safe_float(row, "value", default_value=7.89), 7.89)
+
+
+class TestGetSafeInt(unittest.TestCase):
+    def test_valid_int(self):
+        row = {"value": "42"}
+        self.assertEqual(get_safe_int(row, "value"), 42)
+
+    def test_invalid_int(self):
+        row = {"value": "abc"}
+        self.assertIsNone(get_safe_int(row, "value"))
+
+    def test_missing_key(self):
+        row = {}
+        self.assertIsNone(get_safe_int(row, "value"))
+
+    def test_empty_string(self):
+        row = {"value": ""}
+        self.assertIsNone(get_safe_int(row, "value"))

functions-python/helpers/transform.py

@@ -89,7 +89,7 @@ def to_float(value, default_value: Optional[float] = None) -> Optional[float]:
         return default_value
 
 
-def get_safe_value(row, column_name, default_value=None) -> Optional[str]:
+def get_safe_value(row, column_name, default_value: str = None) -> Optional[str]:
     """
     Get a safe value from the row. If the value is missing or empty, return the default value.
     """
@@ -105,12 +105,23 @@ def get_safe_value(row, column_name, default_value=None) -> Optional[str]:
     return f"{value}".strip()
 
 
-def get_safe_float(row, column_name, default_value=None) -> Optional[float]:
+def get_safe_float(row, column_name, default_value: float = None) -> Optional[float]:
     """
-    Get a safe float value from the row. If the value is missing or cannot be converted to float,
+    Get a safe float value from the row. If the value is missing or cannot be converted to float.
     """
     safe_value = get_safe_value(row, column_name)
     try:
         return float(safe_value)
     except (ValueError, TypeError):
         return default_value
+
+
+def get_safe_int(row, column_name, default_value: int = None) -> Optional[int]:
+    """
+    Get a safe int value from the row. If the value is missing or cannot be converted to int.
+    """
+    safe_value = get_safe_value(row, column_name)
+    try:
+        return int(safe_value)
+    except (ValueError, TypeError):
+        return default_value

functions-python/pmtiles_builder/requirements_dev.txt

@@ -1,4 +1,6 @@
 Faker
 pytest~=7.4.3
 urllib3-mock
-requests-mock
+requests-mock
+psutil
+gcp_storage_emulator

functions-python/pmtiles_builder/src/csv_cache.py

@@ -15,6 +15,8 @@
 #
 import csv
 import os
+from typing import TypedDict, List, Dict
+
 
 from gtfs import stop_txt_is_lat_log_required
 from shared.helpers.logger import get_logger
@@ -28,6 +30,11 @@
 AGENCY_FILE = "agency.txt"
 
 
+class ShapeTrips(TypedDict):
+    shape_id: str
+    trip_ids: List[str]
+
+
 class CsvCache:
     """
     CsvCache provides cached access to GTFS CSV files in a specified working directory.
@@ -49,11 +56,12 @@ def __init__(
         self.workdir = workdir
 
         self.file_data = {}
-        self.trip_to_stops = None
+        self.trip_to_stops: Dict[str, List[str]] = None
         self.route_to_trip = None
-        self.route_to_shape = None
+        self.route_to_shape: Dict[str, Dict[str, ShapeTrips]] = None
         self.stop_to_route = None
         self.stop_to_coordinates = None
+        self.trips_no_shapes_per_route: Dict[str, List[str]] = {}
 
         self.logger.info("Using work directory: %s", self.workdir)
 
@@ -90,41 +98,63 @@ def _read_csv(self, filename) -> list[dict]:
         except Exception as e:
             raise Exception(f"Failed to read CSV file {filename}: {e}") from e
 
-    def get_trip_from_route(self, route_id):
-        if self.route_to_trip is None:
-            self.route_to_trip = {}
-            for row in self.get_file(TRIPS_FILE):
-                route_id = row["route_id"]
-                trip_id = row["trip_id"]
-                if trip_id:
-                    self.route_to_trip.setdefault(route_id, trip_id)
-        return self.route_to_trip.get(route_id, "")
-
-    def get_shape_from_route(self, route_id) -> str:
+    def get_shape_from_route(self, route_id) -> Dict[str, List[ShapeTrips]]:
         """
-        Returns the first shape_id associated with a given route_id from the trips file.
+        Returns a list of shape_ids with associated trip_ids information with a given route_id from the trips file.
         The relationship from the route to the shape is via the trips file.
         Parameters:
             route_id (str): The route identifier to look up.
 
         Returns:
             The corresponding shape id.
+            Example return value: [{'shape_id1': { 'shape_id': 'shape_id1', 'trip_ids': ['trip1', 'trip2']}},
+             {'shape_id': 'shape_id2', 'trip_ids': ['trip3']}}]
         """
         if self.route_to_shape is None:
             self.route_to_shape = {}
             for row in self.get_file(TRIPS_FILE):
-                route_id = row["route_id"]
-                shape_id = row["shape_id"]
-                if shape_id:
-                    self.route_to_shape.setdefault(route_id, shape_id)
-        return self.route_to_shape.get(route_id, "")
+                route_id = get_safe_value(row, "route_id")
+                shape_id = get_safe_value(row, "shape_id")
+                trip_id = get_safe_value(row, "trip_id")
+                if route_id and trip_id:
+                    if shape_id:
+                        route_shapes = self.route_to_shape.setdefault(route_id, {})
+                        shape_trips = route_shapes.setdefault(
+                            shape_id, {"shape_id": shape_id, "trip_ids": []}
+                        )
+                        shape_trips["trip_ids"].append(trip_id)
+                    else:
+                        # Registering the trip without a shape for this route for later retrieval.
+                        trip_no_shapes = (
+                            self.trips_no_shapes_per_route.get(route_id)
+                            if route_id in self.trips_no_shapes_per_route
+                            else None
+                        )
+                        if trip_no_shapes is None:
+                            trip_no_shapes = []
+                            self.trips_no_shapes_per_route[route_id] = trip_no_shapes
+                        trip_no_shapes.append(trip_id)
+        return self.route_to_shape.get(route_id, {})
+
+    def get_trips_without_shape_from_route(self, route_id) -> List[str]:
+        return self.trips_no_shapes_per_route.get(route_id, [])
 
     def get_stops_from_trip(self, trip_id):
-        # Lazy instantiation of the dictionary, because we may not need it al all if there is a shape.
         if self.trip_to_stops is None:
             self.trip_to_stops = {}
             for row in self.get_file(STOP_TIMES_FILE):
-                self.trip_to_stops.setdefault(row["trip_id"], []).append(row["stop_id"])
+                trip_id = get_safe_value(row, "trip_id")
+                stop_id = get_safe_value(row, "stop_id")
+                if trip_id and stop_id:
+                    trip_to_stops = (
+                        self.trip_to_stops.get(trip_id)
+                        if trip_id in self.trip_to_stops
+                        else None
+                    )
+                    if trip_to_stops is None:
+                        trip_to_stops = []
+                        self.trip_to_stops[trip_id] = trip_to_stops
+                    trip_to_stops.append(stop_id)
         return self.trip_to_stops.get(trip_id, [])
 
     def get_coordinates_for_stop(self, stop_id) -> tuple[float, float] | None:

functions-python/pmtiles_builder/src/gtfs_stops_to_geojson.py

@@ -6,7 +6,7 @@
 from csv_cache import CsvCache, ROUTES_FILE, TRIPS_FILE, STOP_TIMES_FILE, STOPS_FILE
 from gtfs import stop_txt_is_lat_log_required
 from shared.helpers.runtime_metrics import track_metrics
-from shared.helpers.transform import get_safe_float
+from shared.helpers.transform import get_safe_float, get_safe_value
 
 logger = logging.getLogger(__name__)
 
@@ -15,7 +15,7 @@ def create_routes_map(routes_data):
     """Creates a dictionary of routes from route data."""
     routes = {}
     for row in routes_data:
-        route_id = row.get("route_id")
+        route_id = get_safe_value(row, "route_id")
         if route_id:
             routes[route_id] = row
     return routes
@@ -26,16 +26,16 @@ def build_stop_to_routes(stop_times_data, trips_data):
     # Build trip_id -> route_id mapping
     trip_to_route = {}
     for row in trips_data:
-        trip_id = row.get("trip_id")
-        route_id = row.get("route_id")
+        trip_id = get_safe_value(row, "trip_id")
+        route_id = get_safe_value(row, "route_id")
         if trip_id and route_id:
             trip_to_route[trip_id] = route_id
 
     # Build stop_id -> set of route_ids
     stop_to_routes = defaultdict(set)
     for row in stop_times_data:
-        trip_id = row.get("trip_id")
-        stop_id = row.get("stop_id")
+        trip_id = get_safe_value(row, "trip_id")
+        stop_id = get_safe_value(row, "stop_id")
         if trip_id and stop_id:
             route_id = trip_to_route.get(trip_id)
             if route_id:
@@ -92,13 +92,13 @@ def convert_stops_to_geojson(csv_cache: CsvCache, output_file):
             },
             "properties": {
                 "stop_id": stop_id,
-                "stop_code": row.get("stop_code", ""),
-                "stop_name": row.get("stop_name", ""),
-                "stop_desc": row.get("stop_desc", ""),
-                "zone_id": row.get("zone_id", ""),
-                "stop_url": row.get("stop_url", ""),
-                "wheelchair_boarding": row.get("wheelchair_boarding", ""),
-                "location_type": row.get("location_type", ""),
+                "stop_code": get_safe_value(row, "stop_code", ""),
+                "stop_name": get_safe_value(row, "stop_name", ""),
+                "stop_desc": get_safe_value(row, "stop_desc", ""),
+                "zone_id": get_safe_value(row, "zone_id", ""),
+                "stop_url": get_safe_value(row, "stop_url", ""),
+                "wheelchair_boarding": get_safe_value(row, "wheelchair_boarding", ""),
+                "location_type": get_safe_value(row, "location_type", ""),
                 "route_ids": route_ids,
                 "route_colors": route_colors,
             },