add + resolve PGM mypy support

Signed-off-by: Martijn Govers <[email protected]>

Author: mgovers

.github/workflows/check-code-quality.yml

@@ -53,3 +53,12 @@ jobs:
         run: |
           pip install pylint .
           pylint power_grid_model_io
+
+      - name: If needed raise error
+        run: |
+
+          if [[ `git status --porcelain --untracked-files=no` ]]; then
+            echo "Formatting not correct! See below the files which need to be reformatted!"
+            git status --porcelain --untracked-files=no
+            exit 1
+          fi

src/power_grid_model_io/converters/pandapower_converter.py

@@ -12,7 +12,15 @@
 import numpy as np
 import pandas as pd
 import structlog
-from power_grid_model import Branch3Side, BranchSide, LoadGenType, WindingType, initialize_array, power_grid_meta_data
+from power_grid_model import (
+    Branch3Side,
+    BranchSide,
+    DatasetType,
+    LoadGenType,
+    WindingType,
+    initialize_array,
+    power_grid_meta_data,
+)
 from power_grid_model.data_types import Dataset, SingleDataset
 
 from power_grid_model_io.converters.base_converter import BaseConverter
@@ -113,10 +121,10 @@ def _serialize_data(self, data: Dataset, extra_info: Optional[ExtraInfo]) -> Pan
             self._extra_info_to_pgm_input_data(extra_info)
 
         # Convert
-        def pgm_output_dtype_checker(check_type: str) -> bool:
+        def pgm_output_dtype_checker(check_type: DatasetType | str) -> bool:
             return all(
                 (
-                    comp_array.dtype == power_grid_meta_data[check_type][component]
+                    comp_array.dtype == power_grid_meta_data[DatasetType[check_type]][component]
                     for component, comp_array in self.pgm_output_data.items()
                 )
             )
@@ -248,7 +256,10 @@ def _extra_info_to_pgm_input_data(self, extra_info: ExtraInfo):  # pylint: disab
         nan = np.iinfo(dtype).min
         all_other_cols = ["i_n"]
         for component, data in self.pgm_output_data.items():
-            input_cols = power_grid_meta_data["input"][component].dtype.names
+            input_cols = power_grid_meta_data[DatasetType.input][component].dtype.names
+            if input_cols is None:
+                input_cols = tuple()
+
             node_cols = [col for col in input_cols if is_node_ref(col)]
             other_cols = [col for col in input_cols if col in all_other_cols]
             if not node_cols + other_cols:
@@ -1594,9 +1605,11 @@ def join_currents(table: str, bus_name: str, i_name: str) -> pd.DataFrame:
         # Combine all branch bus, current and et in one dataframe
         all_i_df = pd.concat(
             [
-                join_currents(table, bus_name, i_name)
-                if not rest_switches_absent[table]
-                else pd.DataFrame(columns=["bus", "element", "et", "i_ka"])
+                (
+                    join_currents(table, bus_name, i_name)
+                    if not rest_switches_absent[table]
+                    else pd.DataFrame(columns=["bus", "element", "et", "i_ka"])
+                )
                 for table, attr_names in switch_attrs.items()
                 for bus_name, i_name in attr_names.items()
             ]

src/power_grid_model_io/converters/pgm_json_converter.py

@@ -151,9 +151,12 @@ def _parse_component(
 
                 # If an attribute doesn't exist, it is added to the extra_info lookup table
                 elif extra_info is not None:
-                    if obj["id"] not in extra_info:
-                        extra_info[obj["id"]] = {}
-                    extra_info[obj["id"]][attribute] = value
+                    obj_id = obj["id"]
+                    if not isinstance(obj_id, int):
+                        raise ValueError(f"Invalid 'id' value for {component} {data_type} data")
+                    if obj_id not in extra_info:
+                        extra_info[obj_id] = {}
+                    extra_info[obj_id][attribute] = value
         return array
 
     def _serialize_data(self, data: Dataset, extra_info: Optional[ExtraInfo]) -> StructuredData:
@@ -271,7 +274,11 @@ def _extract_extra_info(
                 self._extract_extra_component_info(component, entry, reserialized_data, extra_info)
 
     def _extract_extra_component_info(
-        self, component: str, attributes: Dict[str, Any], reserialized_data: SingleDataset, extra_info: ExtraInfo
+        self,
+        component: str,
+        attributes: Dict[str, Any],
+        reserialized_data: Dict[str, List[Dict[str, Any]]],
+        extra_info: ExtraInfo,
     ):
         entry_id = attributes["id"]
         reserialized_entry = self._get_first_by(reserialized_data[component], "id", entry_id)

src/power_grid_model_io/converters/tabular_converter.py

@@ -221,7 +221,7 @@ def _parse_table_filters(self, data: TabularData, table: str, filtering_function
         for filtering_fn in filtering_functions:
             for fn_name, kwargs in filtering_fn.items():
                 fn_ptr = get_function(fn_name)
-                table_mask &= data[table].apply(fn_ptr, axis=1, **kwargs).values
+                table_mask &= cast(pd.DataFrame, data[table]).apply(fn_ptr, axis=1, **kwargs).values
         return table_mask
 
     # pylint: disable = too-many-arguments