CSVMassBalanceOutputProvider accepts Path

src/itzi_core/providers/csv_mass_balance_output.py

@@ -15,6 +15,7 @@
 import csv
 import numbers
 from datetime import datetime
+from pathlib import Path
 
 from itzi_core.data_containers import MassBalanceData
 from itzi_core.providers.base import MassBalanceOutputProvider
@@ -23,30 +24,23 @@
 class CSVMassBalanceOutputProvider(MassBalanceOutputProvider):
     """Writes pre-calculated mass balance data to a CSV file."""
 
-    def __init__(
-        self,
-        file_name: str,
-    ):
-        """Initialize the provider and create the output file with headers."""
+    def __init__(self, file_name: str | Path):
         self.fields = list(MassBalanceData.model_fields.keys())
         self.file_name = self._set_file_name(file_name)
         self._create_file()
 
-    def _set_file_name(self, file_name: str) -> str:
-        """Generate output file name"""
+    def _set_file_name(self, file_name: str | Path) -> Path:
         if not file_name:
             timestamp = datetime.now().strftime("%Y-%m-%dT%H-%M-%S")  # noqa: DTZ005
-            file_name = f"{timestamp}_stats.csv"
-        return file_name
+            return Path(f"{timestamp}_stats.csv")
+        return Path(file_name)
 
     def _create_file(self) -> None:
-        """Create a csv file and write headers"""
-        with open(self.file_name, "w", newline="") as f:
+        with self.file_name.open("w", newline="") as f:
             writer = csv.DictWriter(f, fieldnames=self.fields)
             writer.writeheader()
 
     def log(self, report_data: MassBalanceData) -> None:
-        """Writes a single line of data to the CSV file."""
         line_to_write = {}
 
         for key, value in report_data.model_dump().items():
@@ -57,6 +51,6 @@ def log(self, report_data: MassBalanceData) -> None:
             else:
                 line_to_write[key] = value
 
-        with open(self.file_name, "a", newline="") as f:
+        with self.file_name.open("a", newline="") as f:
             writer = csv.DictWriter(f, fieldnames=self.fields)
             writer.writerow(line_to_write)

tests/test_csv_mass_balance_output.py

@@ -14,6 +14,7 @@
 
 import csv
 from datetime import UTC, datetime, timedelta
+from pathlib import Path
 
 import pytest
 
@@ -24,7 +25,7 @@
 
 @pytest.fixture
 def provider_fixture(tmp_path):
-    return {"file_name": str(tmp_path / "stats.csv")}
+    return {"file_name": tmp_path / "stats.csv"}
 
 
 def test_init_with_custom_filename(provider_fixture):
@@ -35,12 +36,19 @@ def test_init_with_custom_filename(provider_fixture):
     assert provider.file_name == provider_fixture["file_name"]
 
 
+def test_init_with_string_filename(tmp_path):
+    file_name = str(tmp_path / "stats.csv")
+    provider = CSVMassBalanceOutputProvider(file_name=file_name)
+
+    assert provider.file_name == Path(file_name)
+
+
 def test_init_with_default_filename(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)
     provider = CSVMassBalanceOutputProvider(
         file_name="",
     )
-    assert provider.file_name.endswith("_stats.csv")
+    assert provider.file_name.name.endswith("_stats.csv")
 
 
 @pytest.mark.parametrize(
@@ -74,7 +82,7 @@ def test_log_temporal_value(provider_fixture, test_time):
     )
 
     provider.log(test_data)
-    with open(provider_fixture["file_name"], "r") as f:
+    with provider_fixture["file_name"].open() as f:
         reader = csv.DictReader(f)
         assert reader.fieldnames == list(MassBalanceData.model_fields)
         row = next(reader)