a few addition dataset types fixes

Signed-off-by: Santiago Figueroa Manrique <[email protected]>

Author: figueroa1395

src/power_grid_model_io/converters/base_converter.py

@@ -55,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
@@ -72,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
@@ -86,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
@@ -100,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
@@ -114,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.

src/power_grid_model_io/converters/pandapower_converter.py

@@ -86,7 +86,7 @@ def _parse_data(
         Args:
             data: PandaPowerData, 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
 
@@ -103,7 +103,7 @@ def _parse_data(
         self.pp_input_data = data
 
         # Convert
-        if data_type == "input":
+        if data_type == DatasetType.input:
             self._create_input_data()
         else:
             raise ValueError(f"Data type: '{data_type}' is not implemented")
@@ -149,9 +149,9 @@ def pgm_output_dtype_checker(check_type: DatasetType | str) -> bool:
             )
 
         # Convert
-        if pgm_output_dtype_checker("sym_output"):
+        if pgm_output_dtype_checker(DatasetType.sym_output):
             self._create_output_data()
-        elif pgm_output_dtype_checker("asym_output"):
+        elif pgm_output_dtype_checker(DatasetType.asym_output):
             self._create_output_data_3ph()
         else:
             raise TypeError("Invalid output data dictionary supplied.")

src/power_grid_model_io/converters/pgm_json_converter.py

@@ -61,7 +61,8 @@ def _parse_data(self, data: StructuredData, data_type: DatasetType, extra_info:
 
         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:
@@ -70,7 +71,7 @@ def _parse_data(self, data: StructuredData, data_type: DatasetType, extra_info:
 
         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")
@@ -99,7 +100,7 @@ def _parse_dataset(
 
         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:
@@ -108,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 {
@@ -128,7 +129,7 @@ def _parse_component(
           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:

src/power_grid_model_io/converters/tabular_converter.py

@@ -96,7 +96,7 @@ def _parse_data(self, data: TabularData, data_type: DatasetType, extra_info: Opt
         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:
@@ -158,7 +158,7 @@ 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

tests/unit/converters/test_tabular_converter.py

@@ -180,7 +180,7 @@ def test_convert_table_to_component__filters(
     }
     converter._convert_table_to_component(
         data=tabular_data_no_units_no_substitutions,
-        data_type="input",
+        data_type=DatasetType.input,
         table="nodes",
         component=ComponentType.node,
         attributes=node_attributes_with_filter,

tests/validation/test_test_utils.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 import pandas as pd
-from power_grid_model import ComponentType, power_grid_meta_data
+from power_grid_model import ComponentType, DatasetType, power_grid_meta_data
 
 from .utils import component_attributes, extract_extra_info, select_values
 
@@ -17,7 +17,7 @@
 @patch("pathlib.Path.open", mock_open(read_data=MOCK_JSON_DATA))
 def test_component_attributes():
     # Act
-    generator = component_attributes(Path("test.json"), data_type="input")
+    generator = component_attributes(Path("test.json"), data_type=DatasetType.input)
 
     # Assert
     assert list(generator) == [
@@ -30,7 +30,7 @@ def test_component_attributes():
 
 def test_select_values():
     # Arrange
-    input_node_dtype = power_grid_meta_data["input"][ComponentType.node].dtype
+    input_node_dtype = power_grid_meta_data[DatasetType.input][ComponentType.node].dtype
     actual = {ComponentType.node: np.array([(2, 2.0), (4, np.nan), (3, 3.0), (1, 1.0)], dtype=input_node_dtype)}
     expected = {ComponentType.node: np.array([(4, 4.0), (1, np.nan), (3, 3.0)], dtype=input_node_dtype)}
 
@@ -52,7 +52,7 @@ def test_extract_extra_info():
     }
 
     # Act
-    extra_info = extract_extra_info(data=data, data_type="input")
+    extra_info = extract_extra_info(data=data, data_type=DatasetType.input)
 
     # Assert
     assert extra_info[1] == {"name": "bar"}