check on more docs

Signed-off-by: zhen0427 <[email protected]>

Author: zhen0427

tests/unit/converters/test_base_converter.py

@@ -8,7 +8,7 @@
 
 import numpy as np
 import pytest
-from power_grid_model import ComponentType
+from power_grid_model import ComponentType, DatasetType
 
 from power_grid_model_io.converters.base_converter import BaseConverter
 
@@ -53,7 +53,7 @@ def add_extra_info(data, data_type, extra_info):
 
     # Assert
     converter._parse_data.assert_called_once_with(  # type: ignore
-        data={ComponentType.node: [{"id": 1}, {"id": 2}]}, data_type="input", extra_info=ANY
+        data={ComponentType.node: [{"id": 1}, {"id": 2}]}, data_type=DatasetType.input, extra_info=ANY
     )
     assert data == {"foo": 1}
     assert extra_info == {1: "Foo"}
@@ -67,7 +67,7 @@ def test_load_input_data__no_extra_info(converter: DummyConverter):
     data, extra_info = converter.load_input_data(data=mock_data, make_extra_info=False)
 
     # Assert
-    converter._parse_data.assert_called_once_with(data=mock_data, data_type="input", extra_info=None)  # type: ignore
+    converter._parse_data.assert_called_once_with(data=mock_data, data_type=DatasetType.input, extra_info=None)  # type: ignore
     assert extra_info == {}
 
 
@@ -89,7 +89,7 @@ def test_load_update_data(converter: DummyConverter):
 
     # Assert
     converter._parse_data.assert_called_once_with(  # type: ignore
-        data={ComponentType.node: [{"id": 1}, {"id": 2}]}, data_type="update", extra_info=None
+        data={ComponentType.node: [{"id": 1}, {"id": 2}]}, data_type=DatasetType.update, extra_info=None
     )
     assert data == {"foo": 1}
 
@@ -103,7 +103,7 @@ def test_load_sym_output_data(converter: DummyConverter):
 
     # Assert
     converter._parse_data.assert_called_once_with(  # type: ignore
-        data={ComponentType.node: [{"id": 1}, {"id": 2}]}, data_type="sym_output", extra_info=None
+        data={ComponentType.node: [{"id": 1}, {"id": 2}]}, data_type=DatasetType.sym_output, extra_info=None
     )
     assert data == {"foo": 1}
 
@@ -117,7 +117,7 @@ def test_load_asym_output_data(converter: DummyConverter):
 
     # Assert
     converter._parse_data.assert_called_once_with(  # type: ignore
-        data={ComponentType.node: [{"id": 1}, {"id": 2}]}, data_type="asym_output", extra_info=None
+        data={ComponentType.node: [{"id": 1}, {"id": 2}]}, data_type=DatasetType.asym_output, extra_info=None
     )
     assert data == {"foo": 1}
 
@@ -131,7 +131,7 @@ def test_load_sc_output_data(converter: DummyConverter):
 
     # Assert
     converter._parse_data.assert_called_once_with(  # type: ignore
-        data={ComponentType.node: [{"id": 1}, {"id": 2}]}, data_type="sc_output", extra_info=None
+        data={ComponentType.node: [{"id": 1}, {"id": 2}]}, data_type=DatasetType.sc_output, extra_info=None
     )
     assert data == {"foo": 1}
 

tests/unit/converters/test_pandapower_converter_input.py

@@ -115,7 +115,7 @@ def create_input_data():
     create_input_data_mock.side_effect = create_input_data
 
     # Act
-    result = converter._parse_data(data={"bus": pd.DataFrame()}, data_type="input", extra_info=None)
+    result = converter._parse_data(data={"bus": pd.DataFrame()}, data_type=DatasetType.input, extra_info=None)
 
     # Assert
     create_input_data_mock.assert_called_once_with()
@@ -138,7 +138,7 @@ def test_parse_data__extra_info(
     extra_info = MagicMock("extra_info")
 
     # Act
-    converter._parse_data(data={}, data_type="input", extra_info=extra_info)
+    converter._parse_data(data={}, data_type=DatasetType.input, extra_info=extra_info)
 
     # Assert
     create_input_data_mock.assert_called_once_with()
@@ -152,7 +152,7 @@ def test_parse_data__update_data():
 
     # Act/Assert
     with pytest.raises(ValueError):
-        converter._parse_data(data={}, data_type="update", extra_info=None)
+        converter._parse_data(data={}, data_type=DatasetType.update, extra_info=None)
 
 
 def test_fill_pgm_extra_info():
@@ -442,9 +442,9 @@ def test_create_input_data():
     ("create_fn", "table"),
     [
         (PandaPowerConverter._create_pgm_input_nodes, "bus"),
-        (PandaPowerConverter._create_pgm_input_lines, "line"),
+        (PandaPowerConverter._create_pgm_input_lines, ComponentType.line),
         (PandaPowerConverter._create_pgm_input_sources, "ext_grid"),
-        (PandaPowerConverter._create_pgm_input_shunts, "shunt"),
+        (PandaPowerConverter._create_pgm_input_shunts, ComponentType.shunt),
         (PandaPowerConverter._create_pgm_input_sym_gens, "sgen"),
         (PandaPowerConverter._create_pgm_input_sym_loads, "load"),
         (PandaPowerConverter._create_pgm_input_transformers, "trafo"),
@@ -481,7 +481,7 @@ def test_create_pgm_input_nodes(mock_init_array: MagicMock, two_pp_objs: MockDf,
     converter._generate_ids.assert_called_once_with("bus", two_pp_objs.index)
 
     # initialization
-    mock_init_array.assert_called_once_with(data_type="input", component_type=ComponentType.node, shape=2)
+    mock_init_array.assert_called_once_with(data_type=DatasetType.input, component_type=ComponentType.node, shape=2)
 
     # retrieval
     converter._get_pp_attr.assert_any_call("bus", "vn_kv", expected_type="f8")
@@ -514,7 +514,7 @@ def test_create_pgm_input_lines(mock_init_array: MagicMock, two_pp_objs, convert
     converter._get_pgm_ids.assert_any_call("bus", _get_pp_attr("line", "to_bus", expected_type="u4"))
 
     # initialization
-    mock_init_array.assert_called_once_with(data_type="input", component_type=ComponentType.line, shape=2)
+    mock_init_array.assert_called_once_with(data_type=DatasetType.input, component_type=ComponentType.line, shape=2)
 
     # retrieval
     converter._get_pp_attr.assert_any_call("line", "from_bus", expected_type="u4")
@@ -606,7 +606,7 @@ def test_create_pgm_input_sources(mock_init_array: MagicMock, two_pp_objs, conve
     converter._generate_ids.assert_called_once_with("ext_grid", two_pp_objs.index)
 
     # initialization
-    mock_init_array.assert_called_once_with(data_type="input", component_type=ComponentType.source, shape=2)
+    mock_init_array.assert_called_once_with(data_type=DatasetType.input, component_type=ComponentType.source, shape=2)
 
     # retrieval:
     converter._get_pp_attr.assert_any_call("ext_grid", "bus", expected_type="u4")
@@ -671,7 +671,7 @@ def test_create_pgm_input_sym_loads(mock_init_array: MagicMock, two_pp_objs, con
     # administration:
 
     # initialization
-    mock_init_array.assert_called_once_with(data_type="input", component_type=ComponentType.sym_load, shape=3 * 2)
+    mock_init_array.assert_called_once_with(data_type=DatasetType.input, component_type=ComponentType.sym_load, shape=3 * 2)
 
     # retrieval:
     converter._get_pp_attr.assert_any_call("load", "bus", expected_type="u4")
@@ -708,7 +708,7 @@ def test_create_pgm_input_asym_loads(mock_init_array: MagicMock, two_pp_objs, co
     converter._generate_ids.assert_called_once_with("asymmetric_load", two_pp_objs.index)
 
     # initialization
-    mock_init_array.assert_called_once_with(data_type="input", component_type=ComponentType.asym_load, shape=2)
+    mock_init_array.assert_called_once_with(data_type=DatasetType.input, component_type=ComponentType.asym_load, shape=2)
 
     # retrieval:
     converter._get_pp_attr.assert_any_call("asymmetric_load", "bus", expected_type="u4")
@@ -801,7 +801,7 @@ def test_create_pgm_input_shunts(mock_init_array: MagicMock, two_pp_objs, conver
     converter._generate_ids.assert_called_once_with(ComponentType.shunt, two_pp_objs.index)
 
     # initialization
-    mock_init_array.assert_called_once_with(data_type="input", component_type=ComponentType.shunt, shape=2)
+    mock_init_array.assert_called_once_with(data_type=DatasetType.input, component_type=ComponentType.shunt, shape=2)
 
     # retrieval:
     converter._get_pp_attr.assert_any_call(ComponentType.shunt, "bus", expected_type="u4")
@@ -868,7 +868,7 @@ def test_create_pgm_input_transformers(mock_init_array: MagicMock, two_pp_objs,
     )
 
     # initialization
-    mock_init_array.assert_called_once_with(data_type="input", component_type=ComponentType.transformer, shape=2)
+    mock_init_array.assert_called_once_with(data_type=DatasetType.input, component_type=ComponentType.transformer, shape=2)
 
     # retrieval:
     converter._get_pp_attr.assert_any_call("trafo", "hv_bus", expected_type="u4")
@@ -1017,7 +1017,7 @@ def test_create_pgm_input_sym_gens(mock_init_array: MagicMock, two_pp_objs, conv
     converter._generate_ids.assert_called_once_with("sgen", two_pp_objs.index)
 
     # initialization
-    mock_init_array.assert_called_once_with(data_type="input", component_type=ComponentType.sym_gen, shape=2)
+    mock_init_array.assert_called_once_with(data_type=DatasetType.input, component_type=ComponentType.sym_gen, shape=2)
 
     # retrieval:
     converter._get_pp_attr.assert_any_call("sgen", "bus", expected_type="i8")
@@ -1110,7 +1110,7 @@ def test_create_pgm_input_asym_gens(mock_init_array: MagicMock, two_pp_objs, con
     converter._generate_ids.assert_called_once_with("asymmetric_sgen", two_pp_objs.index)
 
     # initialization
-    mock_init_array.assert_called_once_with(data_type="input", component_type=ComponentType.asym_gen, shape=2)
+    mock_init_array.assert_called_once_with(data_type=DatasetType.input, component_type=ComponentType.asym_gen, shape=2)
 
     # retrieval:
     converter._get_pp_attr.assert_any_call("asymmetric_sgen", "bus", expected_type="i8")
@@ -1162,7 +1162,7 @@ def test_create_pgm_input_three_winding_transformers(mock_init_array: MagicMock,
 
     # initialization
     mock_init_array.assert_called_once_with(
-        data_type="input", component_type=ComponentType.three_winding_transformer, shape=2
+        data_type=DatasetType.input, component_type=ComponentType.three_winding_transformer, shape=2
     )
 
     # retrieval:
@@ -1521,7 +1521,7 @@ def test_create_pgm_input_links(mock_init_array: MagicMock, converter):
     pd.testing.assert_index_equal(converter._generate_ids.call_args_list[0].args[1], pd.Index([1, 3]))
 
     # initialization
-    mock_init_array.assert_called_once_with(data_type="input", component_type=ComponentType.link, shape=2)
+    mock_init_array.assert_called_once_with(data_type=DatasetType.input, component_type=ComponentType.link, shape=2)
 
     # assignment:
     pgm: MagicMock = mock_init_array.return_value.__setitem__
@@ -1581,7 +1581,7 @@ def test_create_pgm_input_wards(mock_init_array: MagicMock, two_pp_objs, convert
     # administration:
 
     # initialization
-    mock_init_array.assert_called_once_with(data_type="input", component_type=ComponentType.sym_load, shape=2 * 2)
+    mock_init_array.assert_called_once_with(data_type=DatasetType.input, component_type=ComponentType.sym_load, shape=2 * 2)
 
     # retrieval:
     converter._get_pp_attr.assert_any_call("ward", "bus", expected_type="u4")
@@ -1678,7 +1678,7 @@ def test_create_pgm_input_motors(mock_init_array: MagicMock, two_pp_objs, conver
     converter._generate_ids.assert_called_once_with("motor", two_pp_objs.index, name="motor_load")
 
     # initialization
-    mock_init_array.assert_called_once_with(data_type="input", component_type=ComponentType.sym_load, shape=2)
+    mock_init_array.assert_called_once_with(data_type=DatasetType.input, component_type=ComponentType.sym_load, shape=2)
 
     # retrieval:
     converter._get_pp_attr.assert_any_call("motor", "bus", expected_type="i8")

tests/unit/converters/test_pgm_json_converter.py

@@ -27,7 +27,7 @@ def converter():
 @pytest.fixture
 def structured_input_data():
     input_data = {
-        "node": [
+        ComponentType.node: [
             {"id": 1, "u_rated": 400.0},
             {"id": 2, "u_rated": 400.0, "some_extra_info": 2.1},
         ]
@@ -38,8 +38,8 @@ def structured_input_data():
 @pytest.fixture
 def structured_batch_data():
     batch_data = [
-        {"sym_load": [{"id": 3, "p_specified": 1.0}]},
-        {"sym_load": [{"id": 3, "p_specified": 2.0}, {"id": 4, "p_specified": 3.0}]},
+        {ComponentType.sym_load: [{"id": 3, "p_specified": 1.0}]},
+        {ComponentType.sym_load: [{"id": 3, "p_specified": 2.0}, {"id": 4, "p_specified": 3.0}]},
     ]
     return batch_data
 
@@ -48,13 +48,13 @@ def structured_batch_data():
 def pgm_input_data():
     node = initialize_array(DatasetType.input, ComponentType.node, 2)
     node["id"] = [1, 2]
-    return {"node": node}
+    return {ComponentType.node: node}
 
 
 @pytest.fixture
 def pgm_batch_data():
     line = initialize_array(DatasetType.update, ComponentType.line, (3, 2))
-    return {"line": line}
+    return {ComponentType.line: line}
 
 
 @pytest.fixture
@@ -65,19 +65,19 @@ def pgm_sparse_batch_data():
 
 def test_parse_data(converter: PgmJsonConverter, structured_input_data, structured_batch_data):
     with pytest.raises(PowerGridSerializationError, match="Map or an array expected."):
-        converter._parse_data(data="str", data_type="input", extra_info=None)  # type: ignore
+        converter._parse_data(data="str", data_type=DatasetType.input, extra_info=None)  # type: ignore
 
     # test for input dataset
     extra_info: ExtraInfo = {}
-    pgm_data = converter._parse_data(data=structured_input_data, data_type="input", extra_info=extra_info)
+    pgm_data = converter._parse_data(data=structured_input_data, data_type=DatasetType.input, extra_info=extra_info)
     assert len(pgm_data) == 1
     assert len(pgm_data[ComponentType.node]) == 2
     assert [1, 2] in pgm_data[ComponentType.node]["id"]
     assert [400.0, 400.0] in pgm_data[ComponentType.node]["u_rated"]
     assert extra_info == {2: {"some_extra_info": 2.1}}
 
     # test for batch dataset
-    pgm_batch_data = converter._parse_data(data=structured_batch_data, data_type="update", extra_info=None)
+    pgm_batch_data = converter._parse_data(data=structured_batch_data, data_type=DatasetType.update, extra_info=None)
     assert len(pgm_batch_data) == 1
     assert (pgm_batch_data[ComponentType.sym_load]["indptr"] == np.array([0, 1, 3])).all()
     assert (pgm_batch_data[ComponentType.sym_load]["data"]["id"] == [3, 3, 4]).all()
@@ -86,7 +86,7 @@ def test_parse_data(converter: PgmJsonConverter, structured_input_data, structur
 
 def test_parse_dataset(converter: PgmJsonConverter, structured_input_data):
     extra_info: ExtraInfo = {}
-    pgm_data = converter._parse_dataset(data=structured_input_data, data_type="input", extra_info=extra_info)
+    pgm_data = converter._parse_dataset(data=structured_input_data, data_type=DatasetType.input, extra_info=extra_info)
 
     assert len(pgm_data) == 1
     assert len(pgm_data[ComponentType.node]) == 2
@@ -97,11 +97,11 @@ def test_parse_dataset(converter: PgmJsonConverter, structured_input_data):
 
 def test_parse_component(converter: PgmJsonConverter, structured_input_data):
     objects = list(structured_input_data.values())
-    component = "node"
+    component = ComponentType.node
     extra_info: ExtraInfo = {}
 
     node_array = converter._parse_component(
-        objects=objects[0], component=component, data_type="input", extra_info=extra_info
+        objects=objects[0], component=component, data_type=DatasetType.input, extra_info=extra_info
     )
     assert (len(node_array)) == 2
     assert [1, 2] in node_array["id"]
@@ -111,9 +111,9 @@ def test_parse_component(converter: PgmJsonConverter, structured_input_data):
     node_with_wrong_attr_val = {"id": 3, "u_rated": "fault"}
     objects[0].append(node_with_wrong_attr_val)  # type: ignore
     with pytest.raises(
-        ValueError, match="Invalid 'u_rated' value for node input data: could not convert string to float: 'fault'"
+        ValueError, match="Invalid 'u_rated' value for ComponentType.node DatasetType.input data: could not convert string to float: 'fault'"
     ):
-        converter._parse_component(objects=objects[0], component=component, data_type="input", extra_info=None)
+        converter._parse_component(objects=objects[0], component=component, data_type=DatasetType.input, extra_info=None)
 
 
 def test_serialize_data(converter: PgmJsonConverter, pgm_input_data: SingleDataset, pgm_batch_data: BatchDataset):