Merge pull request #295 from PowerGridModel/feature/use-enum

Code quality improvements: Use new enums of dataset and component type

Author: nitbharambe

src/power_grid_model_io/converters/base_converter.py

@@ -10,6 +10,7 @@
 from typing import Generic, Optional, Tuple, TypeVar
 
 import structlog
+from power_grid_model import DatasetType
 from power_grid_model.data_types import Dataset, SingleDataset
 
 from power_grid_model_io.data_stores.base_data_store import BaseDataStore
@@ -54,7 +55,9 @@ def load_input_data(
 
         data = self._load_data(data)
         extra_info: ExtraInfo = {}
-        parsed_data = self._parse_data(data=data, data_type="input", extra_info=extra_info if make_extra_info else None)
+        parsed_data = self._parse_data(
+            data=data, data_type=DatasetType.input, extra_info=extra_info if make_extra_info else None
+        )
         if isinstance(parsed_data, list):
             raise TypeError("Input data can not be batch data")
         return parsed_data, extra_info
@@ -71,7 +74,7 @@ def load_update_data(self, data: Optional[T] = None) -> Dataset:
 
         """
         data = self._load_data(data)
-        return self._parse_data(data=data, data_type="update", extra_info=None)
+        return self._parse_data(data=data, data_type=DatasetType.update, extra_info=None)
 
     def load_sym_output_data(self, data: Optional[T] = None) -> Dataset:
         """Load symmetric output data
@@ -85,7 +88,7 @@ def load_sym_output_data(self, data: Optional[T] = None) -> Dataset:
 
         """
         data = self._load_data(data)
-        return self._parse_data(data=data, data_type="sym_output", extra_info=None)
+        return self._parse_data(data=data, data_type=DatasetType.sym_output, extra_info=None)
 
     def load_asym_output_data(self, data: Optional[T] = None) -> Dataset:
         """Load asymmetric output data
@@ -99,7 +102,7 @@ def load_asym_output_data(self, data: Optional[T] = None) -> Dataset:
 
         """
         data = self._load_data(data)
-        return self._parse_data(data=data, data_type="asym_output", extra_info=None)
+        return self._parse_data(data=data, data_type=DatasetType.asym_output, extra_info=None)
 
     def load_sc_output_data(self, data: Optional[T] = None) -> Dataset:
         """Load sc output data
@@ -113,7 +116,7 @@ def load_sc_output_data(self, data: Optional[T] = None) -> Dataset:
 
         """
         data = self._load_data(data)
-        return self._parse_data(data=data, data_type="sc_output", extra_info=None)
+        return self._parse_data(data=data, data_type=DatasetType.sc_output, extra_info=None)
 
     def convert(self, data: Dataset, extra_info: Optional[ExtraInfo] = None) -> T:
         """Convert input/update/(a)sym_output data and optionally extra info.
@@ -182,7 +185,7 @@ def _load_data(self, data: Optional[T]) -> T:
         raise ValueError("No data supplied!")
 
     @abstractmethod  # pragma: nocover
-    def _parse_data(self, data: T, data_type: str, extra_info: Optional[ExtraInfo]) -> Dataset:
+    def _parse_data(self, data: T, data_type: DatasetType, extra_info: Optional[ExtraInfo]) -> Dataset:
         pass
 
     @abstractmethod  # pragma: nocover

src/power_grid_model_io/converters/pandapower_converter.py

@@ -8,11 +8,12 @@
 import json
 import logging
 import warnings
+from enum import Enum
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional
 
 import numpy as np
-from power_grid_model import initialize_array
+from power_grid_model import ComponentType, DatasetType, initialize_array
 from power_grid_model.data_types import ComponentList, Dataset, SingleDataset, SinglePythonDataset
 from power_grid_model.utils import json_deserialize, json_serialize
 
@@ -44,23 +45,24 @@ class PgmJsonConverter(BaseConverter[StructuredData]):
 
     def __init__(
         self,
-        source_file: Optional[Union[Path, str]] = None,
-        destination_file: Optional[Union[Path, str]] = None,
+        source_file: Optional[Path | str] = None,
+        destination_file: Optional[Path | str] = None,
         log_level: int = logging.INFO,
     ):
         source = JsonFileStore(file_path=Path(source_file)) if source_file else None
         destination = JsonFileStore(file_path=Path(destination_file)) if destination_file else None
         super().__init__(source=source, destination=destination, log_level=log_level)
 
-    def _parse_data(self, data: StructuredData, data_type: str, extra_info: Optional[ExtraInfo]) -> Dataset:
+    def _parse_data(self, data: StructuredData, data_type: DatasetType, extra_info: Optional[ExtraInfo]) -> Dataset:
         """This function expects Structured data, which can either be a dictionary (single dataset) or a list of
         dictionaries (batch dataset). The structured dataset consists of components + attributes that exist within
         power-grid-model, but can also contain other data. If this data should be saved for later usage an extra_info
         dictionary can be provided when calling this function
 
         Args:
           data: Structured data, which can either be a dictionary or a list of dictionaries
-          data_type: the data type of the dataset, i.e. "input", "update", "sym_output" or "asym_output"
+          data_type: the data type of the dataset, i.e. DatasetType.input, DatasetType.update,
+                     DatasetType.sym_output or DatasetType.asym_output
           extra_info: an optional dictionary where extra component info (that can't be specified in
         power-grid-model data) can be specified
           data: StructuredData:
@@ -69,7 +71,7 @@ def _parse_data(self, data: StructuredData, data_type: str, extra_info: Optional
 
         Returns:
           a dictionary containing the components as keys and their corresponding numpy arrays as values: a
-          power-grid-model "input" or "update" dataset
+          power-grid-model DatasetType.input or DatasetType.update dataset
 
         """
         self._log.debug(f"Loading PGM {data_type} data")
@@ -92,13 +94,13 @@ def _parse_data(self, data: StructuredData, data_type: str, extra_info: Optional
         return result
 
     def _parse_dataset(
-        self, data: SinglePythonDataset, data_type: str, extra_info: Optional[ExtraInfo]
+        self, data: SinglePythonDataset, data_type: DatasetType, extra_info: Optional[ExtraInfo]
     ) -> SingleDataset:
         """This function parses a single Python dataset and returns a power-grid-model input or update dictionary
 
         Args:
           data: a single Python dataset
-          data_type: the data type of the dataset, i.e. "input" or "update"
+          data_type: the data type of the dataset, i.e. DatasetType.input or DatasetType.update
           extra_info: an optional dictionary where extra component info (that can't be specified in
         power-grid-model data) can be specified
           data: SinglePythonDataset:
@@ -107,7 +109,7 @@ def _parse_dataset(
 
         Returns:
           a dictionary containing the components as keys and their corresponding numpy arrays as values: a
-          power-grid-model "input" or "update" dataset
+          power-grid-model DatasetType.input or DatasetType.update dataset
 
         """
         return {
@@ -119,15 +121,15 @@ def _parse_dataset(
 
     @staticmethod
     def _parse_component(
-        objects: ComponentList, component: str, data_type: str, extra_info: Optional[ExtraInfo]
+        objects: ComponentList, component: ComponentType, data_type: DatasetType, extra_info: Optional[ExtraInfo]
     ) -> np.ndarray:
         """This function generates a structured numpy array (power-grid-model native) from a structured dataset
 
         Args:
           objects: a list with dictionaries, where each dictionary contains all attributes of a component
           component: the type of component, eg. node, line, etc. Note: it should be a valid power-grid-model
         component
-          data_type: a string specifying the data type: input/update
+          data_type: a string specifying the data type: DatasetType.input/DatasetType.update
           extra_info: an optional dictionary where extra component info (that can't be specified in
         power-grid-model data) can be specified
           objects: ComponentList:
@@ -152,13 +154,15 @@ def _parse_component(
                     try:
                         array[i][attribute] = value
                     except ValueError as ex:
-                        raise ValueError(f"Invalid '{attribute}' value for {component} {data_type} data: {ex}") from ex
+                        raise ValueError(
+                            f"Invalid '{attribute}' value for {component.value} {data_type.value} data: {ex}"
+                        ) from ex
 
                 # If an attribute doesn't exist, it is added to the extra_info lookup table
                 elif extra_info is not None:
                     obj_id = obj["id"]
                     if not isinstance(obj_id, int):
-                        raise ValueError(f"Invalid 'id' value for {component} {data_type} data")
+                        raise ValueError(f"Invalid 'id' value for {component.value} {data_type.value} data")
                     if obj_id not in extra_info:
                         extra_info[obj_id] = {}
                     extra_info[obj_id][attribute] = value
@@ -215,8 +219,9 @@ def _is_batch(data: Dataset) -> bool:
             is_sparse_batch = isinstance(array, dict) and "indptr" in array and "data" in array
             if is_batch is not None and is_batch != (is_dense_batch or is_sparse_batch):
                 raise ValueError(
-                    f"Mixed {'' if is_batch else 'non-'}batch data "
-                    f"with {'non-' if is_batch else ''}batch data ({component})."
+                    f"Mixed {'' if is_batch else 'non-'}batch data with "
+                    f"{'non-' if is_batch else ''}batch data ("
+                    f"{component.value if isinstance(component, Enum) else component})."
                 )
             is_batch = is_dense_batch or is_sparse_batch
         return bool(is_batch)

src/power_grid_model_io/converters/pgm_json_converter.py

@@ -7,13 +7,14 @@
 
 import inspect
 import logging
+from enum import Enum
 from pathlib import Path
-from typing import Any, Collection, Dict, List, Mapping, Optional, Union, cast
+from typing import Any, Collection, Dict, List, Mapping, Optional, cast
 
 import numpy as np
 import pandas as pd
 import yaml
-from power_grid_model import initialize_array
+from power_grid_model import DatasetType, initialize_array
 from power_grid_model.data_types import Dataset
 
 from power_grid_model_io.converters.base_converter import BaseConverter
@@ -89,13 +90,13 @@ def set_mapping(self, mapping: Mapping[str, Any]) -> None:
         if "multipliers" in mapping:
             self._multipliers = MultiplierMapping(cast(Multipliers, mapping["multipliers"]), logger=self._log)
 
-    def _parse_data(self, data: TabularData, data_type: str, extra_info: Optional[ExtraInfo]) -> Dataset:
+    def _parse_data(self, data: TabularData, data_type: DatasetType, extra_info: Optional[ExtraInfo]) -> Dataset:
         """This function parses tabular data and returns power-grid-model data
 
         Args:
           data: TabularData, i.e. a dictionary with the components as keys and pd.DataFrames as values, with
         attribute names as columns and their values in the table
-          data_type: power-grid-model data type, i.e. "input" or "update"
+          data_type: power-grid-model data type, i.e. DatasetType.input or DatasetType.update
           extra_info: an optional dictionary where extra component info (that can't be specified in
         power-grid-model data) can be specified
           data: TabularData:
@@ -145,9 +146,9 @@ def _parse_data(self, data: TabularData, data_type: str, extra_info: Optional[Ex
     def _convert_table_to_component(  # pylint: disable = too-many-arguments,too-many-positional-arguments
         self,
         data: TabularData,
-        data_type: str,
+        data_type: str | Enum,
         table: str,
-        component: str,
+        component: str | Enum,
         attributes: InstanceAttributes,
         extra_info: Optional[ExtraInfo],
     ) -> Optional[np.ndarray]:
@@ -157,17 +158,17 @@ def _convert_table_to_component(  # pylint: disable = too-many-arguments,too-man
 
         Args:
           data: The full dataset with tabular data
-          data_type: The data type, i.e. "input" or "update"
+          data_type: The data type, i.e. DatasetType.input or DatasetType.update
           table: The name of the table that should be converter
           component: the component for which a power-grid-model array should be made
           attributes: a dictionary with a mapping from the attribute names in the table to the corresponding
         power-grid-model attribute names
           extra_info: an optional dictionary where extra component info (that can't be specified in
         power-grid-model data) can be specified
           data: TabularData:
-          data_type: str:
+          data_type: str | Enum:
           table: str:
-          component: str:
+          component: str | Enum:
           attributes: InstanceAttributes:
           extra_info: Optional[ExtraInfo]:
 
@@ -187,13 +188,16 @@ def _convert_table_to_component(  # pylint: disable = too-many-arguments,too-man
 
         n_records = np.sum(table_mask) if table_mask is not None else len(data[table])
 
+        component_str = component.value if isinstance(component, Enum) else component
+        data_type_str = data_type.value if isinstance(data_type, Enum) else data_type
+
         try:
-            pgm_data = initialize_array(data_type=data_type, component_type=component, shape=n_records)
+            pgm_data = initialize_array(data_type=data_type_str, component_type=component_str, shape=n_records)
         except KeyError as ex:
-            raise KeyError(f"Invalid component type '{component}' or data type '{data_type}'") from ex
+            raise KeyError(f"Invalid component type '{component_str}' or data type '{data_type_str}'") from ex
 
         if "id" not in attributes:
-            raise KeyError(f"No mapping for the attribute 'id' for '{component}s'!")
+            raise KeyError(f"No mapping for the attribute 'id' for '{component_str}s'!")
 
         # Make sure that the "id" column is always parsed first (at least before "extra" is parsed)
         attributes_without_filter = {k: v for k, v in attributes.items() if k != "filters"}
@@ -207,7 +211,7 @@ def _convert_table_to_component(  # pylint: disable = too-many-arguments,too-man
                 data=data,
                 pgm_data=pgm_data,
                 table=table,
-                component=component,
+                component=component_str,
                 attr=attr,
                 col_def=col_def,
                 table_mask=table_mask,
@@ -232,7 +236,7 @@ def _convert_col_def_to_attribute(  # pylint: disable = too-many-arguments,too-m
         data: TabularData,
         pgm_data: np.ndarray,
         table: str,
-        component: str,
+        component: str | Enum,
         attr: str,
         col_def: Any,
         table_mask: Optional[np.ndarray],
@@ -254,7 +258,7 @@ def _convert_col_def_to_attribute(  # pylint: disable = too-many-arguments,too-m
           data: TabularData:
           pgm_data: np.ndarray:
           table: str:
-          component: str:
+          component: str | Enum:
           attr: str:
           col_def: Any:
           extra_info: Optional[ExtraInfo]:
@@ -265,12 +269,15 @@ def _convert_col_def_to_attribute(  # pylint: disable = too-many-arguments,too-m
         """
         # To avoid mistakes, the attributes in the mapping should exist. There is one extra attribute called
         # 'extra' in which extra information can be captured.
+
+        component_str = component.value if isinstance(component, Enum) else component
+
         if pgm_data.dtype.names is None:
-            raise ValueError(f"pgm_data for '{component}s' has no attributes defined. (dtype.names is None)")
+            raise ValueError(f"pgm_data for '{component_str}s' has no attributes defined. (dtype.names is None)")
 
         if attr not in pgm_data.dtype.names and attr not in ["extra", "filters"]:
             attrs = ", ".join(pgm_data.dtype.names)
-            raise KeyError(f"Could not find attribute '{attr}' for '{component}s'. (choose from: {attrs})")
+            raise KeyError(f"Could not find attribute '{attr}' for '{component_str}s'. (choose from: {attrs})")
 
         if attr == "extra":
             # Extra info must be linked to the object IDs, therefore the uuids should be known before extra info can
@@ -421,16 +428,15 @@ def _parse_col_def(  # pylint: disable = too-many-arguments,too-many-positional-
     def _parse_col_def_const(
         data: TabularData,
         table: str,
-        col_def: Union[int, float],
+        col_def: int | float,
         table_mask: Optional[np.ndarray] = None,
     ) -> pd.DataFrame:
         """Create a single column pandas DataFrame containing the const value.
 
         Args:
           data: TabularData:
           table: str:
-          col_def: Union[int:
-          float]:
+          col_def: int | float:
 
         Returns:
 
@@ -602,7 +608,7 @@ def _parse_auto_id(  # pylint: disable = too-many-arguments,too-many-positional-
         table: str,
         ref_table: Optional[str],
         ref_name: Optional[str],
-        key_col_def: Union[str, List[str], Dict[str, str]],
+        key_col_def: str | List[str] | Dict[str, str],
         table_mask: Optional[np.ndarray],
         extra_info: Optional[ExtraInfo],
     ) -> pd.DataFrame:
@@ -845,7 +851,7 @@ def get_id(row: pd.Series) -> int:
 
         return keys.apply(get_id, axis=1).to_list()
 
-    def lookup_id(self, pgm_id: int) -> Dict[str, Union[str, Dict[str, int]]]:
+    def lookup_id(self, pgm_id: int) -> Dict[str, str | Dict[str, int]]:
         """
         Retrieve the original name / key combination of a pgm object
         Args:

src/power_grid_model_io/converters/tabular_converter.py

@@ -8,7 +8,7 @@
 import logging
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Mapping, Optional, Union
+from typing import Any, Mapping, Optional
 
 from power_grid_model_io.converters.tabular_converter import TabularConverter
 from power_grid_model_io.data_stores.base_data_store import LANGUAGE_EN
@@ -37,10 +37,10 @@ class VisionExcelConverter(TabularConverter):
 
     def __init__(  # pylint: disable=too-many-arguments,too-many-positional-arguments
         self,
-        source_file: Optional[Union[Path, str]] = None,
+        source_file: Optional[Path | str] = None,
         language: str = LANGUAGE_EN,
         terms_changed: Optional[dict] = None,
-        mapping_file: Optional[Union[Path, str]] = None,
+        mapping_file: Optional[Path | str] = None,
         log_level: int = logging.INFO,
     ):
         _mapping_file = Path(