From 0351485680e2ada378395eba2234931625e43f45 Mon Sep 17 00:00:00 2001 From: Bram Stoeller Date: Fri, 14 Oct 2022 19:56:36 +0200 Subject: [PATCH] CSV data store (beta) Signed-off-by: Bram Stoeller --- .../data_stores/csv_dir_store.py | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100644 src/power_grid_model_io/data_stores/csv_dir_store.py diff --git a/src/power_grid_model_io/data_stores/csv_dir_store.py b/src/power_grid_model_io/data_stores/csv_dir_store.py new file mode 100644 index 00000000..fa510b94 --- /dev/null +++ b/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 +# +# 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)