Apply some changes from review

Signed-off-by: Thijs Baaijen <[email protected]>

Author: Thijss

src/power_grid_model_ds/_core/model/grids/base.py

@@ -42,7 +42,7 @@
 from power_grid_model_ds._core.model.grids._text_sources import TextSource
 from power_grid_model_ds._core.model.grids.helpers import set_feeder_ids, set_is_feeder
 from power_grid_model_ds._core.utils.pickle import get_pickle_path, load_from_pickle, save_to_pickle
-from power_grid_model_ds._core.utils.serialization import _load_grid_from_json, _save_grid_to_json
+from power_grid_model_ds._core.utils.serialization import load_grid_from_json, save_grid_to_json
 from power_grid_model_ds._core.utils.zip import file2gzip
 
 Self = TypeVar("Self", bound="Grid")
@@ -439,7 +439,7 @@ def from_txt_file(cls, txt_file_path: Path):
             txt_lines = f.readlines()
         return TextSource(grid_class=cls).load_from_txt(*txt_lines)
 
-    def to_json(self, path: Path, **kwargs) -> Path:
+    def serialize(self, path: Path, **kwargs) -> Path:
         """Serialize the grid to JSON format.
 
         Args:
@@ -448,12 +448,12 @@ def to_json(self, path: Path, **kwargs) -> Path:
         Returns:
             Path: The path where the file was saved.
         """
-        return _save_grid_to_json(grid=self, path=path, **kwargs)
+        return save_grid_to_json(grid=self, path=path, **kwargs)
 
     @classmethod
-    def from_json(cls: Type[Self], path: Path) -> Self:
+    def deserialize(cls: Type[Self], path: Path) -> Self:
         """Deserialize the grid from JSON format."""
-        return _load_grid_from_json(path=path, target_grid_class=cls)
+        return load_grid_from_json(path=path, target_grid_class=cls)
 
     def set_feeder_ids(self):
         """Sets feeder and substation id properties in the grids arrays"""

src/power_grid_model_ds/_core/utils/serialization.py

@@ -8,7 +8,7 @@
 import json
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, Type, TypeVar
+from typing import TYPE_CHECKING, Type, TypeVar
 
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
 
@@ -24,43 +24,13 @@
 logger = logging.getLogger(__name__)
 
 
-def _restore_grid_values(grid, json_data: Dict) -> None:
-    """Restore arrays to the grid."""
-    for attr_name, attr_values in json_data.items():
-        if not hasattr(grid, attr_name):
-            continue
-
-        if not issubclass(getattr(grid, attr_name).__class__, FancyArray):
-            expected_type = grid.__dataclass_fields__[attr_name].type
-            cast_value = expected_type(attr_values)
-            setattr(grid, attr_name, cast_value)
-            continue
-
-        try:
-            array_field = grid.find_array_field(getattr(grid, attr_name).__class__)
-            matched_columns = {
-                col: attr_values["data"][col] for col in array_field.type().columns if col in attr_values["data"]
-            }
-            restored_array = array_field.type(**matched_columns)
-            setattr(grid, attr_name, restored_array)
-        except (AttributeError, KeyError, ValueError, TypeError) as e:
-            # Handle restoration failures:
-            # - KeyError: missing "dtype" or "data" keys
-            # - ValueError/TypeError: invalid dtype string or data conversion
-            # - AttributeError: grid methods/attributes missing
-            logger.warning(f"Failed to restore '{attr_name}': {e}")
-
-
-def _save_grid_to_json(
-    grid,
-    path: Path,
-    **kwargs,
-) -> Path:
+def save_grid_to_json(grid, path: Path, strict: bool = True, **kwargs) -> Path:
     """Save a Grid object to JSON format using power-grid-model serialization with extensions support.
 
     Args:
         grid: The Grid object to serialize
         path: The file path to save to
+        strict: Whether to raise an error if the grid object is not serializable.
         **kwargs: Keyword arguments forwarded to json.dump (for example, indent, sort_keys,
             ensure_ascii, etc.).
     Returns:
@@ -69,24 +39,25 @@ def _save_grid_to_json(
     path.parent.mkdir(parents=True, exist_ok=True)
 
     serialized_data = {}
+
     for field in dataclasses.fields(grid):
         if field.name in ["graphs", "_id_counter"]:
             continue
 
         field_value = getattr(grid, field.name)
-        if isinstance(field_value, (int, float, str, bool)):
-            serialized_data[field.name] = field_value
-            continue
-
-        if not isinstance(field_value, FancyArray):
-            raise NotImplementedError(f"Serialization for field of type '{type(field_value)}' is not implemented.")
 
-        if field_value.size == 0:
+        if isinstance(field_value, FancyArray):
+            serialized_data[field.name] = {
+                "data": {name: field_value[name].tolist() for name in field_value.dtype.names},
+            }
             continue
 
-        serialized_data[field.name] = {
-            "data": {name: field_value[name].tolist() for name in field_value.dtype.names},
-        }
+        try:
+            json.dumps(field_value)
+        except TypeError as e:
+            if strict:
+                raise
+            logger.warning(f"Failed to serialize '{field.name}': {e}")
 
     # Write to file
     with open(path, "w", encoding="utf-8") as f:
@@ -95,7 +66,7 @@ def _save_grid_to_json(
     return path
 
 
-def _load_grid_from_json(path: Path, target_grid_class: Type[G]) -> G:
+def load_grid_from_json(path: Path, target_grid_class: Type[G]) -> G:
     """Load a Grid object from JSON format with cross-type loading support.
 
     Args:
@@ -108,7 +79,30 @@ def _load_grid_from_json(path: Path, target_grid_class: Type[G]) -> G:
     with open(path, "r", encoding="utf-8") as f:
         input_data = json.load(f)
 
-    target_grid = target_grid_class.empty()
-    _restore_grid_values(target_grid, input_data)
+    grid = target_grid_class.empty()
+    _restore_grid_values(grid, input_data)
+    graph_class = grid.graphs.__class__
+    grid.graphs = graph_class.from_arrays(grid)
+    return grid
+
+
+def _restore_grid_values(grid: G, json_data: dict) -> None:
+    """Restore arrays to the grid."""
+    for attr_name, attr_values in json_data.items():
+        if not hasattr(grid, attr_name):
+            logger.warning(f"Unexpected attribute '{attr_name}'")
+            continue
+
+        grid_attr = getattr(grid, attr_name)
+        attr_class = grid_attr.__class__
+        if isinstance(grid_attr, FancyArray):
+            if extra := set(attr_values["data"]) - set(grid_attr.columns):
+                logger.warning(f"{attr_name} has extra columns: {extra}")
+
+            matched_columns = {col: attr_values["data"][col] for col in grid_attr.columns if col in attr_values["data"]}
+            restored_array = attr_class(**matched_columns)
+            setattr(grid, attr_name, restored_array)
+            continue
 
-    return target_grid
+        # load other values
+        setattr(grid, attr_name, attr_class(attr_values))

tests/integration/visualizer_tests.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: MPL-2.0
 
 from dataclasses import dataclass
+from pathlib import Path
 
 from power_grid_model_ds import Grid
 from power_grid_model_ds._core.visualizer.app import visualize
@@ -50,6 +51,11 @@ def visualize_grid_with_links():
 
 
 if __name__ == "__main__":
+    r_grid = get_radial_grid()
+    r_grid.serialize(Path("json_path"))
+
+    new_grid = Grid.deserialize(Path("json_path"))
+
     visualize_grid()
     # visualize_coordinated_grid()
     # visualize_grid_with_links()

tests/unit/utils/test_serialization.py

@@ -14,8 +14,8 @@
 from power_grid_model_ds import Grid
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
 from power_grid_model_ds._core.utils.serialization import (
-    _load_grid_from_json,
-    _save_grid_to_json,
+    load_grid_from_json,
+    save_grid_to_json,
 )
 from power_grid_model_ds.arrays import LineArray
 from power_grid_model_ds.arrays import NodeArray as BaseNodeArray
@@ -87,11 +87,11 @@ class TestSerializationRoundtrips:
     def test_basic_serialization_roundtrip(self, basic_grid: Grid, tmp_path: Path):
         """Test basic serialization roundtrip for all formats"""
         path = tmp_path / "test.json"
-        result_path = _save_grid_to_json(basic_grid, path)
+        result_path = save_grid_to_json(basic_grid, path)
         assert result_path.exists()
 
         # Load and verify
-        loaded_grid = _load_grid_from_json(path, target_grid_class=Grid)
+        loaded_grid = load_grid_from_json(path, target_grid_class=Grid)
         array_equal(loaded_grid.node, basic_grid.node)
         array_equal(loaded_grid.line, basic_grid.line)
         assert list(loaded_grid.node.id) == list(basic_grid.node.id)
@@ -100,8 +100,8 @@ def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, tmp
         """Test extended serialization preserving custom data"""
         path = tmp_path / "extended.json"
 
-        _save_grid_to_json(extended_grid, path)
-        loaded_grid = _load_grid_from_json(path, target_grid_class=ExtendedGrid)
+        save_grid_to_json(extended_grid, path)
+        loaded_grid = load_grid_from_json(path, target_grid_class=ExtendedGrid)
 
         # Verify core data
         assert loaded_grid.node.size == extended_grid.node.size
@@ -120,10 +120,10 @@ def test_empty_grid_handling(self, tmp_path: Path):
         json_path = tmp_path / "empty.json"
 
         # Should handle empty grids
-        _save_grid_to_json(empty_grid, json_path)
+        save_grid_to_json(empty_grid, json_path)
 
         # Should load back as empty
-        loaded_json = _load_grid_from_json(json_path, target_grid_class=Grid)
+        loaded_json = load_grid_from_json(json_path, target_grid_class=Grid)
         assert loaded_json.node.size == 0
 
 
@@ -135,8 +135,8 @@ def test_basic_to_extended_loading(self, basic_grid: Grid, tmp_path: Path):
         path = tmp_path / "basic.json"
 
         # Save basic grid
-        _save_grid_to_json(basic_grid, path)
-        loaded_grid = _load_grid_from_json(path, target_grid_class=ExtendedGrid)
+        save_grid_to_json(basic_grid, path)
+        loaded_grid = load_grid_from_json(path, target_grid_class=ExtendedGrid)
 
         # Core data should transfer
         array_equal(loaded_grid.node, basic_grid.node)
@@ -147,8 +147,8 @@ def test_extended_to_basic_loading(self, extended_grid: ExtendedGrid, tmp_path:
         path = tmp_path / "extended.json"
 
         # Save extended grid
-        _save_grid_to_json(extended_grid, path)
-        loaded_grid = _load_grid_from_json(path, target_grid_class=Grid)
+        save_grid_to_json(extended_grid, path)
+        loaded_grid = load_grid_from_json(path, target_grid_class=Grid)
 
         # Core data should transfer
         array_equal(loaded_grid.node, extended_grid.node)
@@ -189,10 +189,10 @@ class GridWithCustomArray(Grid):
 
         # Test JSON serialization
         json_path = tmp_path / "custom_array.json"
-        _save_grid_to_json(grid, json_path)
+        save_grid_to_json(grid, json_path)
 
         # Load back and verify
-        loaded_grid = _load_grid_from_json(json_path, target_grid_class=GridWithCustomArray)
+        loaded_grid = load_grid_from_json(json_path, target_grid_class=GridWithCustomArray)
 
         # Verify core data
         assert loaded_grid.node.size == 2