Create separate function for inject_extra_info and handle batch data

Signed-off-by: Bram Stoeller <[email protected]>

Author: bramstoeller

src/power_grid_model/manual_testing.py

@@ -8,7 +8,7 @@
 
 import json
 from pathlib import Path
-from typing import Any, Dict, List, IO, Optional, Union
+from typing import IO, Any, Dict, List, Optional, Union
 
 import numpy as np
 
@@ -256,7 +256,7 @@ def export_json_data(
     data: Union[Dict[str, np.ndarray], List[Dict[str, np.ndarray]]],
     indent: Optional[int] = 2,
     compact: bool = False,
-    extra_info: Optional[Dict[int, Any]] = None,
+    extra_info: Optional[Union[Dict[int, Any], List[Dict[int, Any]]]] = None,
 ):
     """
     export json data
@@ -272,13 +272,8 @@ def export_json_data(
         Save to file
     """
     json_data = convert_numpy_to_python(data)
-
-    # Inject extra info
     if extra_info is not None:
-        for component, objects in json_data.items():
-            for obj in objects:
-                if obj["id"] in extra_info:
-                    obj["extra"] = extra_info[obj["id"]]
+        _inject_extra_info(data=json_data, extra_info=extra_info)
 
     with open(json_file, mode="w", encoding="utf-8") as file_pointer:
         if compact and indent:
@@ -289,6 +284,38 @@ def export_json_data(
             json.dump(json_data, file_pointer, indent=indent)
 
 
+def _inject_extra_info(
+    data: Union[Dict[str, List[Dict[str, Union[float, int]]]], List[Dict[str, List[Dict[str, Union[float, int]]]]]],
+    extra_info: Union[Dict[int, Any], List[Dict[int, Any]]],
+):
+    """
+    Injects extra info to the objects by ID
+
+    Args:
+        data: Power Grid Model Python data, as written to pgm json files.
+        extra_info: A dictionary indexed by object id. The value may be anything.
+
+    """
+    if isinstance(data, list):
+        if isinstance(extra_info, list):
+            # If both data and extra_info are lists, expect one extra info set per batch
+            for batch, info in zip(data, extra_info):
+                _inject_extra_info(batch, info)
+        else:
+            # If only data is a list, copy extra_info for each batch
+            for batch in data:
+                _inject_extra_info(batch, extra_info)
+    elif isinstance(data, dict):
+        if not isinstance(extra_info, dict):
+            raise TypeError("Invalid extra info data type")
+        for component, objects in data.items():
+            for obj in objects:
+                if obj["id"] in extra_info:
+                    obj["extra"] = extra_info[obj["id"]]
+    else:
+        raise TypeError("Invalid data type")
+
+
 def _compact_json_dump(data: Any, io_stream: IO[str], indent: int, max_level: int, level: int = 0):
     """Custom compact JSON writer that is intended to put data belonging to a single object on a single line.
 

tests/unit/test_manual_testing.py

@@ -4,17 +4,19 @@
 
 import io
 from pathlib import Path
-from unittest.mock import patch, mock_open, MagicMock
+from unittest.mock import MagicMock, mock_open, patch
 
 import numpy as np
 import pytest
+
 from power_grid_model.manual_testing import (
+    _compact_json_dump,
+    _inject_extra_info,
     convert_batch_to_list_data,
     convert_numpy_to_python,
     convert_python_to_numpy,
     export_json_data,
     is_nan,
-    _compact_json_dump,
 )
 
 
@@ -136,6 +138,73 @@ def test_export_json_data(convert_mock: MagicMock, open_mock: MagicMock, json_du
     json_dump_mock.assert_called_once_with({"foo": [{"val": 123}]}, open_mock(), indent=2)
 
 
+@patch("json.dump")
+@patch("builtins.open", new_callable=mock_open)
+@patch("power_grid_model.manual_testing.convert_numpy_to_python")
+@patch("power_grid_model.manual_testing._inject_extra_info")
+def test_export_json_data_extra_info(
+    extra_info_mock: MagicMock, convert_mock: MagicMock, _open_mock: MagicMock, _json_dump_mock: MagicMock
+):
+    convert_mock.return_value = {"foo": [{"id": 123}]}
+    export_json_data(json_file=Path(), data={}, extra_info={123: "Extra information"})
+    extra_info_mock.assert_called_once_with(data={"foo": [{"id": 123}]}, extra_info={123: "Extra information"})
+
+
+def test_inject_extra_info_single():
+    data = {"node": [{"id": 0, "foo": 123}, {"id": 1, "bar": 456}], "line": [{"id": 2, "baz": 789}]}
+    extra_info = {2: 42, 1: {"sheet": "Nodes", "Number": "00123"}}
+    _inject_extra_info(data=data, extra_info=extra_info)
+    assert data == {
+        "node": [{"id": 0, "foo": 123}, {"id": 1, "bar": 456, "extra": {"sheet": "Nodes", "Number": "00123"}}],
+        "line": [{"id": 2, "baz": 789, "extra": 42}],
+    }
+
+
+def test_inject_extra_info_batch():
+    data = [
+        {"node": [{"id": 0, "foo": 111}, {"id": 1, "bar": 222}], "line": [{"id": 2, "baz": 333}]},
+        {"node": [{"id": 0, "foo": 444}, {"id": 1, "bar": 555}], "line": [{"id": 2, "baz": 666}]},
+    ]
+    extra_info = [{2: 42, 1: {"sheet": "Nodes", "Number": "00123"}}, {2: 43, 0: None}]
+    _inject_extra_info(data=data, extra_info=extra_info)
+    assert data == [
+        {
+            "node": [{"id": 0, "foo": 111}, {"id": 1, "bar": 222, "extra": {"sheet": "Nodes", "Number": "00123"}}],
+            "line": [{"id": 2, "baz": 333, "extra": 42}],
+        },
+        {
+            "node": [{"id": 0, "foo": 444, "extra": None}, {"id": 1, "bar": 555}],
+            "line": [{"id": 2, "baz": 666, "extra": 43}],
+        },
+    ]
+
+
+def test_inject_extra_info_batch_copy_info():
+    data = [
+        {"node": [{"id": 0, "foo": 111}, {"id": 1, "bar": 222}], "line": [{"id": 2, "baz": 333}]},
+        {"node": [{"id": 0, "foo": 444}, {"id": 1, "bar": 555}], "line": [{"id": 2, "baz": 666}]},
+    ]
+    extra_info = {2: 42, 1: {"sheet": "Nodes", "Number": "00123"}}
+    _inject_extra_info(data=data, extra_info=extra_info)
+    assert data == [
+        {
+            "node": [{"id": 0, "foo": 111}, {"id": 1, "bar": 222, "extra": {"sheet": "Nodes", "Number": "00123"}}],
+            "line": [{"id": 2, "baz": 333, "extra": 42}],
+        },
+        {
+            "node": [{"id": 0, "foo": 444}, {"id": 1, "bar": 555, "extra": {"sheet": "Nodes", "Number": "00123"}}],
+            "line": [{"id": 2, "baz": 666, "extra": 42}],
+        },
+    ]
+
+
+def test_inject_extra_info_single_dataset_with_batch_info():
+    data = {"node": [{"id": 0, "foo": 123}, {"id": 1, "bar": 456}], "line": [{"id": 2, "baz": 789}]}
+    extra_info = [{2: 42, 1: {"sheet": "Nodes", "Number": "00123"}}, {2: 43, 0: None}]
+    with pytest.raises(TypeError):
+        _inject_extra_info(data=data, extra_info=extra_info)
+
+
 def test_compact_json_dump():
     data = {
         "node": [{"id": 1, "x": 2}, {"id": 3, "x": 4}],