simplified: id generation to kick in only when string conversion fails

nagacharan-tangirala · Dec 27, 2024 · f5aaa7d · f5aaa7d
1 parent f3afca1
commit f5aaa7d
Showing 1 changed file with 9 additions and 6 deletions.
diff --git a/src/prep_disolv/vehicle/sumo.py b/src/prep_disolv/vehicle/sumo.py
@@ -92,12 +92,15 @@ def get_parquet_file(self) -> Path:
 
     def get_vehicle_id_from_pool(self, vehicle_id_str: str) -> int:
         """Return a vehicle ID from the map."""
-        if vehicle_id_str in self.vehicle_id_pool.keys():
-            vehicle_id = self.vehicle_id_pool[vehicle_id_str]
-        else:
-            vehicle_id = self.vehicle_id_init
-            self.vehicle_id_pool[vehicle_id_str] = vehicle_id
-            self.vehicle_id_init = self.vehicle_id_init + 1
+        try:
+            vehicle_id = int(vehicle_id_str)
+        except ValueError:
+            if vehicle_id_str in self.vehicle_id_pool.keys():
+                vehicle_id = self.vehicle_id_pool[vehicle_id_str]
+            else:
+                vehicle_id = self.vehicle_id_init
+                self.vehicle_id_pool[vehicle_id_str] = vehicle_id
+                self.vehicle_id_init = self.vehicle_id_init + 1
         return vehicle_id
 
     def _convert_fcd_to_parquet(self) -> None: