Update generation script to fix issues and make it faster

3 files changed, 50 insertions, 16 deletions

diff --git a/src/gtfs_vigo_stops/src/report_writer.py b/src/gtfs_vigo_stops/src/report_writer.py
index 9560895..695931f 100644
--- a/src/gtfs_vigo_stops/src/report_writer.py
+++ b/src/gtfs_vigo_stops/src/report_writer.py
@@ -7,8 +7,6 @@ import json
 import os
 from typing import Any, Dict, List
 
-from pyproj import Transformer
-
 from src.logger import get_logger
 from src.proto.stop_schedule_pb2 import Epsg25829, StopArrivals
 
@@ -18,8 +16,8 @@ def write_stop_protobuf(
     date: str,
     stop_code: str,
     arrivals: List[Dict[str, Any]],
-    stop_lat: float,
-    stop_lon: float,
+    stop_x: float,
+    stop_y: float,
 ) -> None:
     """
     Write stop arrivals data to a Protobuf file.
@@ -32,12 +30,9 @@ def write_stop_protobuf(
     """
     logger = get_logger("report_writer")
 
-    transformer = Transformer.from_crs(4326, 25829, always_xy=True)
-    x, y = transformer.transform(stop_lon, stop_lat)
-
     item = StopArrivals(
         stop_id=stop_code,
-        location=Epsg25829(x=x, y=y),
+        location=Epsg25829(x=stop_x, y=stop_y),
         arrivals=[
             StopArrivals.ScheduledArrival(
                 service_id=arrival["service_id"],
diff --git a/src/gtfs_vigo_stops/src/stops.py b/src/gtfs_vigo_stops/src/stops.py
index 2c58387..44a6d2d 100644
--- a/src/gtfs_vigo_stops/src/stops.py
+++ b/src/gtfs_vigo_stops/src/stops.py
@@ -3,6 +3,8 @@ import os
 from dataclasses import dataclass
 from typing import Dict, Optional
 
+from pyproj import Transformer
+
 from src.logger import get_logger
 
 logger = get_logger("stops")
@@ -16,17 +18,34 @@ class Stop:
     stop_lat: Optional[float]
     stop_lon: Optional[float]
 
+    stop_25829_x: Optional[float] = None
+    stop_25829_y: Optional[float] = None
+
 
 CACHED_STOPS: dict[str, dict[str, Stop]] = {}
+CACHED_BY_CODE: dict[str, dict[str, Stop]] = {}
 
 
 def get_all_stops_by_code(feed_dir: str) -> Dict[str, Stop]:
+    if feed_dir in CACHED_BY_CODE:
+        return CACHED_BY_CODE[feed_dir]
+
+    transformer = Transformer.from_crs(4326, 25829, always_xy=True)
+
     stops_by_code: Dict[str, Stop] = {}
     all_stops = get_all_stops(feed_dir)
 
     for stop in all_stops.values():
         if stop.stop_code:
-            stops_by_code[stop.stop_code] = stop
+            stop_25829_x, stop_25829_y = transformer.transform(
+                stop.stop_lon, stop.stop_lat
+            )
+            stop.stop_25829_x = stop_25829_x
+            stop.stop_25829_y = stop_25829_y
+
+            stops_by_code[get_numeric_code(stop.stop_code)] = stop
+
+    CACHED_BY_CODE[feed_dir] = stops_by_code
 
     return stops_by_code
 
diff --git a/src/gtfs_vigo_stops/stop_report.py b/src/gtfs_vigo_stops/stop_report.py
index 7db751c..ab6bac8 100644
--- a/src/gtfs_vigo_stops/stop_report.py
+++ b/src/gtfs_vigo_stops/stop_report.py
@@ -2,6 +2,7 @@ import argparse
 import os
 import shutil
 import sys
+import time
 import traceback
 from typing import Any, Dict, List
 
@@ -256,9 +257,24 @@ def process_date(
             logger.warning(f"No stop arrivals found for date {date}")
             return date, {}
 
+        logger.info(
+            f"Writing stop reports for {len(stop_arrivals)} stops for date {date}"
+        )
+
+        # Write individual stop JSON files
+        writing_start_time = time.perf_counter()
+        for stop_code, arrivals in stop_arrivals.items():
+            write_stop_json(output_dir, date, stop_code, arrivals)
+        writing_end_time = time.perf_counter()
+        writing_elapsed = writing_end_time - writing_start_time
+
+        logger.info(
+            f"Finished writing stop JSON reports for date {date} in {writing_elapsed:.2f}s"
+        )
+
         # Write individual stop JSON files
+        writing_start_time = time.perf_counter()
         for stop_code, arrivals in stop_arrivals.items():
-            # Get the stop from 'stops' by value to get the coords
             stop_by_code = stops_by_code.get(stop_code)
 
             if stop_by_code is not None:
@@ -267,18 +283,22 @@ def process_date(
                     date,
                     stop_code,
                     arrivals,
-                    stop_by_code.stop_lat or 0.0,
-                    stop_by_code.stop_lon or 0.0,
+                    stop_by_code.stop_25829_x or 0.0,
+                    stop_by_code.stop_25829_y or 0.0,
                 )
 
-            write_stop_json(output_dir, date, stop_code, arrivals)
+        writing_end_time = time.perf_counter()
+        writing_elapsed = writing_end_time - writing_start_time
+
+        logger.info(
+            f"Finished writing stop protobuf reports for date {date} in {writing_elapsed:.2f}s"
+        )
+
+        logger.info(f"Processed {len(stop_arrivals)} stops for date {date}")
 
-        # Create summary for index
         stop_summary = {
             stop_code: len(arrivals) for stop_code, arrivals in stop_arrivals.items()
         }
-        logger.info(f"Processed {len(stop_arrivals)} stops for date {date}")
-
         return date, stop_summary
     except Exception as e:
         logger.error(f"Error processing date {date}: {e}")