CSV data store (beta)

Signed-off-by: Bram Stoeller <bram.stoeller@alliander.com>

Author: bramstoeller

src/power_grid_model_io/data_stores/csv_dir_store.py

@@ -0,0 +1,48 @@
+# SPDX-FileCopyrightText: 2022 Contributors to the Power Grid Model IO project <dynamic.grid.calculation@alliander.com>
+#
+# SPDX-License-Identifier: MPL-2.0
+"""
+CSV Directory Store
+"""
+
+from pathlib import Path
+from typing import Any, Dict, List
+
+import pandas as pd
+
+from power_grid_model_io.data_stores.base_data_store import BaseDataStore
+from power_grid_model_io.data_types import TabularData
+
+
+class CsvDirStore(BaseDataStore[TabularData]):
+    """
+    CSV Directory Store
+
+    The first row of each .csv file is expected to contain the column names, unless specified differently by an
+    extension of this class.
+    """
+
+    __slots__ = ("_dir_path", "_csv_kwargs", "_header_rows")
+
+    def __init__(self, dir_path: Path, **csv_kwargs):
+        super().__init__()
+        self._dir_path = dir_path
+        self._csv_kwargs: Dict[str, Any] = csv_kwargs
+        self._header_rows: List[int] = [0]
+
+    def load(self) -> TabularData:
+        """
+        Load all CSV files in a directory as tabular data.
+        """
+        data: Dict[str, pd.DataFrame] = {}
+        for path in self._dir_path.glob("*.csv"):
+            data[path.stem] = pd.read_csv(filepath_or_buffer=path, header=self._header_rows, **self._csv_kwargs)
+
+        return TabularData(**data)
+
+    def save(self, data: TabularData) -> None:
+        """
+        Store each table in data as a separate CSV file
+        """
+        for table_name, table_data in data.items():
+            table_data.to_csv(path_or_buf=self._dir_path / f"{table_name}.csv", **self._csv_kwargs)