Update unit tests. Fix get_stops_from_trip

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/src/csv_cache.py

@@ -30,7 +30,7 @@
 AGENCY_FILE = "agency.txt"
 
 
-class Shapes(TypedDict):
+class ShapeTrips(TypedDict):
     shape_id: str
     trip_ids: List[str]
 
@@ -56,12 +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: Dict[str, Dict[str, List[str]]] = 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 = Dict[str, List[str]]
+        self.trips_no_shapes_per_route: Dict[str, List[str]] = {}
 
         self.logger.info("Using work directory: %s", self.workdir)
 
@@ -111,7 +111,7 @@ def get_trip_from_route(self, route_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) -> List[ShapeTrips]:
         """
         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.
@@ -120,38 +120,63 @@ def get_shape_from_route(self, route_id) -> str:
 
         Returns:
             The corresponding shape id.
-            Example return value: [{'shape_id1': ['trip1', 'trip2']}, {'shape_id2': ['trip3']}]
+            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"]
-                trip_id = row["trip_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.get(route_id, {})
+                        route_shapes = self.route_to_shape.get(route_id, None)
+                        if route_shapes is None:
+                            route_shapes = {}
+                            self.route_to_shape[route_id] = route_shapes
                         if shape_id not in route_shapes:
-                            route_shapes[shape_id] = []
-                        route_shapes[shape_id].append(trip_id)
-                        self.route_to_shape[route_id] = route_shapes
+                            shape_trips = {"shape_id": shape_id, "trip_ids": []}
+                            route_shapes[shape_id] = shape_trips
+                        else:
+                            shape_trips = route_shapes[shape_id]
+                        shape_trips["trip_ids"].append(trip_id)
                     else:
-                        trip_no_shapes = self.trip_to_stops.get(route_id, [])
+                        # 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)
-                        self.trips_no_shapes[route_id] = trip_no_shapes
         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[route_id] if route_id in self.trips_no_shapes else []
+            self.trips_no_shapes_per_route[route_id]
+            if route_id in self.trips_no_shapes_per_route
+            else []
         )
 
     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

@@ -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:

functions-python/pmtiles_builder/src/main.py

@@ -39,13 +39,14 @@
     STOPS_FILE,
     AGENCY_FILE,
     SHAPES_FILE,
+    ShapeTrips,
 )
 from gtfs_stops_to_geojson import convert_stops_to_geojson
 from shared.database_gen.sqlacodegen_models import Gtfsdataset, Gtfsfeed
 from shared.helpers.logger import get_logger, init_logger
 from shared.helpers.runtime_metrics import track_metrics
 from shared.database.database import with_db_session
-from shared.helpers.transform import get_safe_value
+from shared.helpers.transform import get_safe_value, get_safe_float
 
 init_logger()
 
@@ -349,7 +350,12 @@ def _create_shapes_index(self) -> dict:
                     if not line:
                         break
                     row = dict(zip(columns, next(csv.reader([line]))))
-                    sid = row["shape_id"]
+                    sid = get_safe_value(row, "shape_id")
+                    if not sid:
+                        self.logger.warning(
+                            "Missing shape_id at line %s, skipping.", row
+                        )
+                        continue
                     shapes_index.setdefault(sid, []).append(pos)
                     count += 1
                     if count % 1000000 == 0:
@@ -371,11 +377,21 @@ def _get_shape_points(self, shape_id, index):
                     f.seek(pos)
                     line = f.readline()
                     row = dict(zip(index["columns"], next(csv.reader([line]))))
+                    shape_pt_lon = get_safe_float(row, "shape_pt_lon")
+                    shape_pt_lat = get_safe_float(row, "shape_pt_lat")
+                    shape_pt_sequence = get_safe_float(row, "shape_pt_sequence", 0)
+                    if shape_pt_lon is None or shape_pt_lat is None:
+                        self.logger.warning(
+                            "Invalid coordinates for shape_id %s at position %d, skipping.",
+                            shape_id,
+                            pos,
+                        )
+                        continue
                     points.append(
                         (
-                            float(row["shape_pt_lon"]),
-                            float(row["shape_pt_lat"]),
-                            int(row["shape_pt_sequence"]),
+                            shape_pt_lon,
+                            shape_pt_lat,
+                            shape_pt_sequence,
                         )
                     )
             points.sort(key=lambda x: x[2])
@@ -444,7 +460,7 @@ def _create_routes_geojson(self):
                             "Processed route %d (route_id: %s)", i, route_id
                         )
 
-                # geojson_file.write("\n]}")
+                geojson_file.write("\n]}")
 
             if missing_coordinates_routes:
                 self.logger.info(
@@ -457,7 +473,7 @@ def _create_routes_geojson(self):
             raise Exception(f"Failed to create routes GeoJSON: {e}") from e
 
     def get_route_coordinates(self, route_id, shapes_index) -> List[RouteCoordinates]:
-        shapes: Dict[str, List[str]] = self.csv_cache.get_shape_from_route(route_id)
+        shapes: Dict[str, ShapeTrips] = self.csv_cache.get_shape_from_route(route_id)
         result: List[RouteCoordinates] = []
         if shapes:
             for shape_id, trip_ids in shapes.items():
@@ -545,7 +561,6 @@ def _create_routes_json(self):
             routes = []
             for row in self.csv_cache.get_file(ROUTES_FILE):
                 route_id = get_safe_value(row, "route_id", "")
-                shape_ids = self.csv_cache.get_shape_from_route(route_id)
                 route = {
                     "routeId": route_id,
                     "routeName": get_safe_value(row, "route_long_name", "")
@@ -554,7 +569,6 @@ def _create_routes_json(self):
                     "color": f"#{get_safe_value(row, 'route_color', '000000')}",
                     "textColor": f"#{get_safe_value(row, 'route_text_color', 'FFFFFF')}",
                     "routeType": f"{get_safe_value(row, 'route_type', '')}",
-                    "shapes_ids": shape_ids,
                 }
                 routes.append(route)
 

functions-python/pmtiles_builder/src/scripts/pmtiles_builder_verifier.py

@@ -34,7 +34,7 @@ def download_feed_files(feed_stable_id: str, dataset_stable_id: str):
         logging.info(f"Downloading {url}")
         filename = f"{dataset_stable_id}/extracted/{file}"
         try:
-            download_to_local(feed_stable_id, url, filename, True)
+            download_to_local(feed_stable_id, url, filename, False)
         except Exception as e:
             logging.warning(f"Failed to download {file}: {e}")