Merge pull request #1042 from PowerGridModel/feature/clean-up-remaining-enums-2

Use of enumerations: Clean remaining enums in other places

Author: nitbharambe

src/power_grid_model/data_types.py

@@ -101,16 +101,20 @@
 """
 A batch dataset is a dictionary where the keys are the component types and the values are :class:`BatchComponentData`
 
-- Example: {"node": :class:`DenseBatchArray`, "line": :class:`SparseBatchArray`,
-            "link": :class:`DenseBatchColumnarData`, "transformer": :class:`SparseBatchColumnarData`}
+- Example: {
+    ComponentType.node: :class:`DenseBatchArray`,
+    ComponentType.line: :class:`SparseBatchArray`,
+    ComponentType.link: :class:`DenseBatchColumnarData`,
+    ComponentType.transformer: :class:`SparseBatchColumnarData`
+  }
 """
 
 BatchList = _BatchList
 """
 A batch list is an alternative representation of a batch. It is a list of single datasets, where each single dataset
 is actually a batch. The batch list is intended as an intermediate data type, during conversions.
 
-- Example: [:class:`SingleDataset`, {"node": :class:`SingleDataset`}]
+- Example: [:class:`SingleDataset`, {ComponentType.node: :class:`SingleDataset`}]
 """
 
 BatchPythonDataset = _BatchPythonDataset
@@ -121,8 +125,8 @@
 
 - Example:
 
-  [{"line": [{"id": 3, "from_status": 0, "to_status": 0, ...}],},
-   {"line": [{"id": 3, "from_status": 1, "to_status": 1, ...}],}]
+  [{ComponentType.line: [{"id": 3, "from_status": 0, "to_status": 0, ...}],},
+   {ComponentType.line: [{"id": 3, "from_status": 1, "to_status": 1, ...}],}]
 """
 
 ColumnarData = _ColumnarData
@@ -162,10 +166,14 @@
 
 - Examples:
 
-    - single: {"node": :class:`SingleArray`, "line": :class:`SingleColumnarData`}
+    - single: {ComponentType.node: :class:`SingleArray`, ComponentType.line: :class:`SingleColumnarData`}
 
-    - batch: {"node": :class:`DenseBatchArray`, "line": :class:`SparseBatchArray`,
-             "link": :class:`DenseBatchColumnarData`, "transformer": :class:`SparseBatchColumnarData`}
+    - batch: {
+        ComponentType.node: :class:`DenseBatchArray`,
+        ComponentType.line: :class:`SparseBatchArray`,
+        ComponentType.link: :class:`DenseBatchColumnarData`,
+        ComponentType.transformer: :class:`SparseBatchColumnarData`
+      }
 
 """
 
@@ -216,14 +224,14 @@
   - single:
 
     {
-      "node": [{"id": 1, "u_rated": 10500.0}, {"id": 2, "u_rated": 10500.0}],
-      "line": [{"id": 3, "from_node": 1, "to_node": 2, ...}],
+      ComponentType.node: [{"id": 1, "u_rated": 10500.0}, {"id": 2, "u_rated": 10500.0}],
+      ComponentType.line: [{"id": 3, "from_node": 1, "to_node": 2, ...}],
     }
 
   - batch:
 
-    [{"line": [{"id": 3, "from_status": 0, "to_status": 0, ...}],},
-     {"line": [{"id": 3, "from_status": 1, "to_status": 1, ...}],}]
+    [{ComponentType.line: [{"id": 3, "from_status": 0, "to_status": 0, ...}],},
+     {ComponentType.line: [{"id": 3, "from_status": 1, "to_status": 1, ...}],}]
 """
 
 RealValue = _RealValue
@@ -240,7 +248,10 @@
 - Examples:
 
     - structure: <1d-array>
-    - concrete: array([(0, 10500.0), (0, 10500.0)], dtype=power_grid_meta_data["input"]["node"].dtype)
+    - concrete: array(
+        [(0, 10500.0), (0, 10500.0)],
+        dtype=power_grid_meta_data[DatasetType.input][ComponentType.node].dtype
+      )
 """
 
 SingleColumn = _SingleColumn
@@ -253,8 +264,11 @@
     - structure: <1d-array>
     - concrete:
 
-        - array([0, 1], dtype=power_grid_meta_data["input"]["node"].dtype.fields["id"][0])
-        - array([10500.0, 10500.0], dtype=power_grid_meta_data["input"]["node"].dtype.fields["u_rated"][0])
+        - array([0, 1], dtype=power_grid_meta_data[DatasetType.input][ComponentType.node].dtype.fields["id"][0])
+        - array(
+            [10500.0, 10500.0],
+            dtype=power_grid_meta_data[DatasetType.input][ComponentType.node].dtype.fields["u_rated"][0]
+          )
 """
 
 
@@ -277,7 +291,7 @@
 A single dataset is a dictionary where the keys are the component types and the values are
 :class:`ComponentData`
 
-- Example: {"node": :class:`SingleArray`, "line": :class:`SingleColumnarData`}
+- Example: {ComponentType.node: :class:`SingleArray`, ComponentType.line: :class:`SingleColumnarData`}
 """
 
 SinglePythonDataset = _SinglePythonDataset
@@ -289,8 +303,8 @@
 - Example:
 
   {
-    "node": [{"id": 1, "u_rated": 10500.0}, {"id": 2, "u_rated": 10500.0}], 
-    "line": [{"id": 3, "from_node": 1, "to_node": 2, ...}],
+    ComponentType.node: [{"id": 1, "u_rated": 10500.0}, {"id": 2, "u_rated": 10500.0}],
+    ComponentType.line: [{"id": 3, "from_node": 1, "to_node": 2, ...}],
   }
 """
 

tests/unit/test_0Z_model_validation.py

@@ -9,6 +9,7 @@
 import numpy as np
 import pytest
 
+from power_grid_model._core.dataset_definitions import ComponentType
 from power_grid_model._core.utils import convert_batch_dataset_to_batch_list
 from power_grid_model.enum import TapChangingStrategy
 
@@ -129,12 +130,12 @@ def test_single_validation(
         assert np.all(input_array["id"][indexer_array] == ids_array)
 
     # test calculate with only node and source result
-    kwargs = dict(base_kwargs, **{"output_component_types": ["node", "source"]})
+    kwargs = dict(base_kwargs, **{"output_component_types": [ComponentType.node, ComponentType.source]})
     result = calculation_function(model, **supported_kwargs(kwargs=kwargs, supported=calculation_args))
-    assert set(result.keys()) == {"node", "source"}
-    kwargs = dict(base_kwargs, **{"output_component_types": {"node", "source"}})
+    assert set(result.keys()) == {ComponentType.node, ComponentType.source}
+    kwargs = dict(base_kwargs, **{"output_component_types": {ComponentType.node, ComponentType.source}})
     result = calculation_function(model, **supported_kwargs(kwargs=kwargs, supported=calculation_args))
-    assert set(result.keys()) == {"node", "source"}
+    assert set(result.keys()) == {ComponentType.node, ComponentType.source}
 
 
 @pytest.mark.parametrize(

tests/unit/test_dataset.py

@@ -54,8 +54,10 @@ def test_const_dataset__conflicting_data():
     with pytest.raises(PowerGridError):
         CConstDataset(
             data={
-                ComponentType.node: np.zeros(1, dtype=power_grid_meta_data["input"][ComponentType.node]),
-                "sym_load": np.zeros(1, dtype=power_grid_meta_data["update"]["sym_load"]),
+                ComponentType.node: np.zeros(1, dtype=power_grid_meta_data[DatasetType.input][ComponentType.node]),
+                ComponentType.sym_load: np.zeros(
+                    1, dtype=power_grid_meta_data[DatasetType.update][ComponentType.sym_load]
+                ),
             }
         )
 
@@ -112,14 +114,16 @@ def test_const_dataset__sparse_batch_data(dataset_type):
             "indptr": np.array([0, 2, 3, 3]),
         },
         ComponentType.sym_load: {
-            "data": np.zeros(shape=2, dtype=power_grid_meta_data[dataset_type]["sym_load"]),
+            "data": np.zeros(shape=2, dtype=power_grid_meta_data[dataset_type][ComponentType.sym_load]),
             "indptr": np.array([0, 0, 1, 2]),
         },
         ComponentType.asym_load: {
-            "data": np.zeros(shape=4, dtype=power_grid_meta_data[dataset_type]["asym_load"]),
+            "data": np.zeros(shape=4, dtype=power_grid_meta_data[dataset_type][ComponentType.asym_load]),
             "indptr": np.array([0, 2, 3, 4]),
         },
-        ComponentType.link: np.zeros(shape=(batch_size, 4), dtype=power_grid_meta_data[dataset_type]["link"]),
+        ComponentType.link: np.zeros(
+            shape=(batch_size, 4), dtype=power_grid_meta_data[dataset_type][ComponentType.link]
+        ),
     }
 
     dataset = CConstDataset(data, dataset_type)