allow CSV write to receive None. Update type hint accordingly (#207)

Author: lrntct

src/itzi/providers/base.py

@@ -12,6 +12,8 @@
 GNU General Public License for more details.
 """
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from typing import Mapping, TYPE_CHECKING, Union, Tuple
 
@@ -36,14 +38,14 @@ def get_origin(self) -> Tuple[float, float]:
         return (domain_data.north, domain_data.west)
 
     @abstractmethod
-    def get_domain_data(self) -> "DomainData":
+    def get_domain_data(self) -> DomainData:
         """Return a DomainData object."""
         pass
 
     @abstractmethod
     def get_array(
-        self, map_key: str, current_time: "datetime"
-    ) -> Tuple["np.ndarray", "datetime", "datetime"]:
+        self, map_key: str, current_time: datetime
+    ) -> Tuple[np.ndarray, datetime, datetime]:
         """Take a given map key and current time
         return a numpy array associated with its start and end time
         if no map is found, return None instead of an array
@@ -61,13 +63,13 @@ def __init__(self, config: Mapping) -> None:
 
     @abstractmethod
     def write_arrays(
-        self, array_dict: Mapping[str, "np.ndarray"], sim_time: Union["datetime", "timedelta"]
+        self, array_dict: Mapping[str, np.ndarray], sim_time: Union[datetime, timedelta]
     ) -> None:
         """Write all arrays for the current time step."""
         pass
 
     @abstractmethod
-    def finalize(self, final_data: "SimulationData") -> None:
+    def finalize(self, final_data: SimulationData) -> None:
         """Finalize outputs and cleanup."""
         pass
 
@@ -82,12 +84,12 @@ def __init__(self, config: Mapping) -> None:
 
     @abstractmethod
     def write_vector(
-        self, drainage_data: "DrainageNetworkData", sim_time: Union["datetime", "timedelta"]
+        self, drainage_data: DrainageNetworkData | None, sim_time: Union[datetime, timedelta]
     ) -> None:
         """Write simulation data for current time step."""
         pass
 
     @abstractmethod
-    def finalize(self, drainage_data: "DrainageNetworkData") -> None:
+    def finalize(self, drainage_data: DrainageNetworkData | None) -> None:
         """Finalize outputs and cleanup."""
         pass

src/itzi/providers/csv_output.py

@@ -13,6 +13,7 @@
 """
 
 from __future__ import annotations
+
 from datetime import datetime, timedelta
 from typing import TypedDict, TYPE_CHECKING, Tuple, List
 from io import StringIO
@@ -88,9 +89,11 @@ def __init__(self, config: CSVVectorOutputConfig) -> None:
         print(self.existing_max_time)
 
     def write_vector(
-        self, drainage_data: DrainageNetworkData, sim_time: datetime | timedelta
+        self, drainage_data: DrainageNetworkData | None, sim_time: datetime | timedelta
     ) -> None:
         """Save simulation data for current time step."""
+        if not drainage_data:
+            return
         # Validate time on first write
         self._validate_time_on_first_write(sim_time)
         # Convert sim_time to ISO8601 format
@@ -106,7 +109,7 @@ def write_vector(
         self._update_csv(sim_time_str, "link", drainage_data.links)
         self.number_of_writes["link"] += 1
 
-    def finalize(self, drainage_data: DrainageNetworkData) -> None:
+    def finalize(self, drainage_data: DrainageNetworkData | None) -> None:
         """Finalize outputs and cleanup."""
         pass
 

src/itzi/providers/grass_output.py

@@ -12,6 +12,8 @@
 GNU General Public License for more details.
 """
 
+from __future__ import annotations
+
 from typing import Dict, TypedDict, Union, TYPE_CHECKING
 
 import numpy as np
@@ -51,7 +53,7 @@ def __init__(self, config: GrassRasterOutputConfig) -> None:
         self.output_maplist = {k: [] for k in self.out_map_names.keys()}
 
     def _write_array(
-        self, array: np.ndarray, map_key: str, sim_time: Union["datetime", "timedelta"]
+        self, array: np.ndarray, map_key: str, sim_time: Union[datetime, timedelta]
     ) -> None:
         """Write simulation data for current time step."""
         suffix = str(self.record_counter[map_key]).zfill(4)
@@ -66,7 +68,7 @@ def _write_array(
         self.record_counter[map_key] += 1
 
     def write_arrays(
-        self, array_dict: Dict[str, np.ndarray], sim_time: Union["datetime", "timedelta"]
+        self, array_dict: Dict[str, np.ndarray], sim_time: Union[datetime, timedelta]
     ) -> None:
         for arr_key, arr in array_dict.items():
             if isinstance(arr, np.ndarray):
@@ -76,7 +78,7 @@ def _write_max_array(self, arr_max, map_key):
         map_max_name = f"{self.out_map_names[map_key]}_max"
         self.grass_interface.write_raster_map(arr_max, map_max_name, map_key, hmin=0.0)
 
-    def finalize(self, final_data: "SimulationData") -> None:
+    def finalize(self, final_data: SimulationData) -> None:
         """Finalize outputs and cleanup."""
 
         # Write the final raster maps
@@ -109,7 +111,7 @@ def __init__(self, config: GrassVectorOutputConfig) -> None:
         self.vector_drainage_maplist = []
 
     def write_vector(
-        self, drainage_data: "DrainageNetworkData", sim_time: Union["datetime", "timedelta"]
+        self, drainage_data: DrainageNetworkData | None, sim_time: Union[datetime, timedelta]
     ) -> None:
         """Write drainage simulation data for current time step."""
         if self.drainage_map_name and drainage_data:
@@ -122,7 +124,7 @@ def write_vector(
             self.vector_drainage_maplist.append((map_name, sim_time))
             self.record_counter += 1
 
-    def finalize(self, drainage_data: "DrainageNetworkData") -> None:
+    def finalize(self, drainage_data: DrainageNetworkData | None) -> None:
         """Finalize outputs and cleanup."""
         if self.drainage_map_name and drainage_data:
             self.grass_interface.register_maps_in_stds(