From 9c3a16119d695f2f905be9ab16ea6c3f7b81537d Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Mon, 13 Oct 2025 18:55:41 +0200 Subject: [PATCH 01/19] feat: setup grid serialization methods Signed-off-by: jaapschoutenalliander --- src/power_grid_model_ds/__init__.py | 16 +- .../_core/model/grids/base.py | 5 +- .../_core/utils/serialization.py | 364 ++++++++++++ tests/unit/utils/test_serialization.py | 548 ++++++++++++++++++ 4 files changed, 931 insertions(+), 2 deletions(-) create mode 100644 src/power_grid_model_ds/_core/utils/serialization.py create mode 100644 tests/unit/utils/test_serialization.py diff --git a/src/power_grid_model_ds/__init__.py b/src/power_grid_model_ds/__init__.py index 78190b0..4f5c6aa 100644 --- a/src/power_grid_model_ds/__init__.py +++ b/src/power_grid_model_ds/__init__.py @@ -5,5 +5,19 @@ from power_grid_model_ds._core.load_flow import PowerGridModelInterface from power_grid_model_ds._core.model.graphs.container import GraphContainer from power_grid_model_ds._core.model.grids.base import Grid +from power_grid_model_ds._core.utils.serialization import ( + load_grid_from_json, + load_grid_from_msgpack, + save_grid_to_json, + save_grid_to_msgpack, +) -__all__ = ["Grid", "GraphContainer", "PowerGridModelInterface"] +__all__ = [ + "Grid", + "GraphContainer", + "PowerGridModelInterface", + "save_grid_to_json", + "save_grid_to_msgpack", + "load_grid_from_json", + "load_grid_from_msgpack", +] diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py index 403af22..35838f5 100644 --- a/src/power_grid_model_ds/_core/model/grids/base.py +++ b/src/power_grid_model_ds/_core/model/grids/base.py @@ -360,7 +360,10 @@ def get_downstream_nodes(self, node_id: int, inclusive: bool = False): ) def cache(self, cache_dir: Path, cache_name: str, compress: bool = True): - """Cache Grid to a folder + """Cache Grid to a folder using pickle format. + + Note: Consider using save_to_json() or save_to_msgpack() for better + interoperability and standardized format. Args: cache_dir (Path): The directory to save the cache to. diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py new file mode 100644 index 0000000..0c8de93 --- /dev/null +++ b/src/power_grid_model_ds/_core/utils/serialization.py @@ -0,0 +1,364 @@ +# SPDX-FileCopyrightText: Contributors to the Power Grid Model project +# +# SPDX-License-Identifier: MPL-2.0 + +"""Serialization utilities for Grid objects using power-grid-model serialization with extensions support.""" + +import dataclasses +import json +import logging +from ast import literal_eval +from pathlib import Path +from typing import Dict, Optional + +import msgpack +import numpy as np +from power_grid_model.utils import json_deserialize, json_serialize, msgpack_deserialize, msgpack_serialize + +from power_grid_model_ds._core.load_flow import PGM_ARRAYS, PowerGridModelInterface +from power_grid_model_ds._core.model.arrays.base.array import FancyArray + +# Constants +EXTENDED_COLUMNS_KEY = "extended_columns" +CUSTOM_ARRAYS_KEY = "custom_arrays" +EXTENSIONS_KEY = "pgm_ds_extensions" + +logger = logging.getLogger(__name__) + + +def _extract_extensions_data(grid) -> Dict[str, Dict]: + """Extract extended columns and non-PGM arrays from a Grid object. + + Args: + grid: The Grid object + + Returns: + Dict containing extensions data with keys EXTENDED_COLUMNS_KEY and CUSTOM_ARRAYS_KEY + """ + extensions: dict = {EXTENDED_COLUMNS_KEY: {}, CUSTOM_ARRAYS_KEY: {}} + + for field in dataclasses.fields(grid): + if field.name in ["graphs", "_id_counter"]: + continue + + array = getattr(grid, field.name) + if not isinstance(array, FancyArray) or array.size == 0: + continue + + array_name = field.name + + if array_name in PGM_ARRAYS: + # Extract extended columns for PGM arrays + _extract_extended_columns(grid, array_name, array, extensions) + else: + # Store custom arrays not in PGM_ARRAYS + extensions[CUSTOM_ARRAYS_KEY][array_name] = {"dtype": str(array.dtype), "data": array.data.tolist()} + + return extensions + + +def _extract_extended_columns(grid, array_name: str, array: FancyArray, extensions: Dict) -> None: + """Extract extended columns from a PGM array.""" + try: + interface = PowerGridModelInterface(grid=grid) + # pylint: disable=protected-access # Accessing internal method for extension extraction + pgm_array = interface._create_power_grid_array(array_name) + pgm_columns = set(pgm_array.dtype.names or []) + ds_columns = set(array.columns) + + # Find extended columns (columns in DS but not in PGM) + extended_cols = ds_columns - pgm_columns + if extended_cols: + extensions[EXTENDED_COLUMNS_KEY][array_name] = {col: array[col].tolist() for col in extended_cols} + except (AttributeError, KeyError, TypeError, ValueError) as e: + # Handle various failure modes: + # - KeyError: array_name not found in PGM arrays + # - AttributeError: array missing dtype/columns or interface method missing + # - TypeError/ValueError: invalid array configuration or data conversion issues + logger.warning(f"Failed to extract extensions for array '{array_name}': {e}") + extensions[CUSTOM_ARRAYS_KEY][array_name] = {"dtype": str(array.dtype), "data": array.data.tolist()} + + +def _restore_extensions_data(grid, extensions_data: Dict) -> None: + """Restore extended columns and custom arrays to a Grid object. + + Args: + grid: The Grid object to restore extensions to + extensions_data: Extensions data from _extract_extensions_data + """ + # Restore extended columns + _restore_extended_columns(grid, extensions_data.get(EXTENDED_COLUMNS_KEY, {})) + + # Restore custom arrays + _restore_custom_arrays(grid, extensions_data.get(CUSTOM_ARRAYS_KEY, {})) + + +def _restore_extended_columns(grid, extended_columns: Dict) -> None: + """Restore extended columns to existing arrays.""" + for array_name, extended_cols in extended_columns.items(): + if not hasattr(grid, array_name): + logger.warning(f"Grid has no attribute '{array_name}' to restore") + continue + + array = getattr(grid, array_name) + if not isinstance(array, FancyArray) or array.size == 0: + continue + + for col_name, values in extended_cols.items(): + # if hasattr(array, col_name): + try: + array[col_name] = values + except (AttributeError, IndexError, ValueError, TypeError) as e: + # Handle assignment failures: + # - IndexError: array size mismatch + # - ValueError/TypeError: incompatible data types + # - AttributeError: array doesn't support assignment + logger.warning(f"Failed to restore column '{col_name}' in array '{array_name}': {e}") + + +def _parse_dtype(dtype_str: str) -> np.dtype: + """Parse a dtype string into a numpy dtype.""" + if not isinstance(dtype_str, str): + raise ValueError(f"Invalid dtype string: {dtype_str}") + + # Use numpy's dtype parsing - handle both eval-style and direct strings + if dtype_str.startswith("dtype("): + clean_dtype_str = dtype_str.replace("dtype(", "").replace(")", "") + else: + clean_dtype_str = dtype_str + + # Use eval for complex dtype strings like "[('field', 'type'), ...]" + if clean_dtype_str.startswith("[") and clean_dtype_str.endswith("]"): + return np.dtype(literal_eval(clean_dtype_str)) + return np.dtype(clean_dtype_str) + + +def _construct_numpy_from_list(raw_data, dtype: np.dtype) -> np.ndarray: + """Construct a numpy array from a list with the specified dtype.""" + if dtype.names: # Structured dtype + # Convert from list of lists to list of tuples for structured array + if isinstance(raw_data[0], (list, tuple)) and len(raw_data[0]) == len(dtype.names): + data = np.array([tuple(row) for row in raw_data], dtype=dtype) + else: + data = np.array(raw_data, dtype=dtype) + else: + data = np.array(raw_data, dtype=dtype) + return data + + +def _restore_custom_arrays(grid, custom_arrays: Dict) -> None: + """Restore custom arrays to the grid.""" + for array_name, array_info in custom_arrays.items(): + if not hasattr(grid, array_name): + continue + + try: + dtype = _parse_dtype(dtype_str=array_info["dtype"]) + data = _construct_numpy_from_list(array_info["data"], dtype) + array_field = grid.find_array_field(getattr(grid, array_name).__class__) + restored_array = array_field.type(data=data) + setattr(grid, array_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 custom array '{array_name}': {e}") + + +def _create_grid_from_input_data(input_data: Dict, target_grid_class=None): + """Create a Grid object from power-grid-model input data. + + Args: + input_data: Power-grid-model input data + target_grid_class: Optional Grid class to create. If None, uses default Grid. + + Returns: + Grid object populated with the input data + """ + if target_grid_class is not None: + # Create empty grid of target type and populate it with input data + target_grid = target_grid_class.empty() + interface = PowerGridModelInterface(grid=target_grid, input_data=input_data) + return interface.create_grid_from_input_data() + + # Use default Grid type + interface = PowerGridModelInterface(input_data=input_data) + return interface.create_grid_from_input_data() + + +def _extract_msgpack_data(data: bytes, **kwargs): + """Extract input data and extensions from MessagePack data.""" + try: + data_dict = msgpack.unpackb(data, raw=False) + if isinstance(data_dict, dict) and EXTENSIONS_KEY in data_dict: + # Extract extensions and deserialize core data + extensions = data_dict.pop(EXTENSIONS_KEY, {}) + core_data = msgpack.packb(data_dict) + input_data = msgpack_deserialize(core_data, **kwargs) + else: + # No extensions, use power-grid-model directly + input_data = msgpack_deserialize(data, **kwargs) + extensions = {EXTENDED_COLUMNS_KEY: {}, CUSTOM_ARRAYS_KEY: {}} + except (msgpack.exceptions.ExtraData, ValueError, TypeError) as e: + # Handle MessagePack parsing failures: + # - ExtraData: malformed MessagePack data + # - ValueError/TypeError: invalid data structure or type issues + logger.warning(f"Failed to extract extensions from MessagePack data: {e}") + input_data = msgpack_deserialize(data, **kwargs) + extensions = {EXTENDED_COLUMNS_KEY: {}, CUSTOM_ARRAYS_KEY: {}} + + return input_data, extensions + + +def _get_serialization_path(path: Path, format_type: str = "auto") -> Path: + """Get the correct path for serialization format. + + Args: + path: Base path + format_type: "json", "msgpack", or "auto" to detect from extension + + Returns: + Path: Path with correct extension + """ + if format_type == "auto": + if path.suffix.lower() in [".json"]: + format_type = "json" + elif path.suffix.lower() in [".msgpack", ".mp"]: + format_type = "msgpack" + else: + # Default to JSON + format_type = "json" + + if format_type == "json" and path.suffix.lower() != ".json": + return path.with_suffix(".json") + if format_type == "msgpack" and path.suffix.lower() not in [".msgpack", ".mp"]: + return path.with_suffix(".msgpack") + + return path + + +def save_grid_to_json( + grid, + path: Path, + use_compact_list: bool = True, + indent: Optional[int] = None, + preserve_extensions: bool = True, +) -> 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 + use_compact_list: Whether to use compact list format + indent: JSON indentation (None for compact, positive int for indentation) + preserve_extensions: Whether to save extended columns and custom arrays + Returns: + Path: The path where the file was saved + """ + path.parent.mkdir(parents=True, exist_ok=True) + + # Convert Grid to power-grid-model input format and serialize + interface = PowerGridModelInterface(grid=grid) + input_data = interface.create_input_from_grid() + + core_data = json_serialize(input_data, use_compact_list=use_compact_list) + + # Parse and add extensions if requested + serialized_data = json.loads(core_data) + if preserve_extensions: + extensions = _extract_extensions_data(grid) + if extensions[EXTENDED_COLUMNS_KEY] or extensions[CUSTOM_ARRAYS_KEY]: + serialized_data[EXTENSIONS_KEY] = extensions + + # Write to file + with open(path, "w", encoding="utf-8") as f: + json.dump(serialized_data, f, indent=indent if indent and indent > 0 else None) + + return path + + +def load_grid_from_json(path: Path, target_grid_class=None): + """Load a Grid object from JSON format with cross-type loading support. + + Args: + path: The file path to load from + target_grid_class: Optional Grid class to load into. If None, uses default Grid. + + Returns: + Grid: The deserialized Grid object of the specified target class + """ + with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + + # Extract extensions and deserialize core data + extensions = data.pop(EXTENSIONS_KEY, {EXTENDED_COLUMNS_KEY: {}, CUSTOM_ARRAYS_KEY: {}}) + input_data = json_deserialize(json.dumps(data)) + + # Create grid and restore extensions + grid = _create_grid_from_input_data(input_data, target_grid_class) + _restore_extensions_data(grid, extensions) + + return grid + + +def save_grid_to_msgpack(grid, path: Path, use_compact_list: bool = True, preserve_extensions: bool = True) -> Path: + """Save a Grid object to MessagePack format with extensions support. + + Args: + grid: The Grid object to serialize + path: The file path to save to + use_compact_list: Whether to use compact list format + preserve_extensions: Whether to save extended columns and custom arrays + + Returns: + Path: The path where the file was saved + """ + path.parent.mkdir(parents=True, exist_ok=True) + + # Convert Grid to power-grid-model input format and serialize + interface = PowerGridModelInterface(grid=grid) + input_data = interface.create_input_from_grid() + + core_data = msgpack_serialize(input_data, use_compact_list=use_compact_list) + + # Add extensions if requested (requires re-serialization for MessagePack) + if preserve_extensions: + extensions = _extract_extensions_data(grid) + if extensions[EXTENDED_COLUMNS_KEY] or extensions[CUSTOM_ARRAYS_KEY]: + core_dict = msgpack.unpackb(core_data, raw=False) + core_dict[EXTENSIONS_KEY] = extensions + serialized_data = msgpack.packb(core_dict) + else: + serialized_data = core_data + else: + serialized_data = core_data + + # Write to file + with open(path, "wb") as f: + f.write(serialized_data) + + return path + + +def load_grid_from_msgpack(path: Path, target_grid_class=None): + """Load a Grid object from MessagePack format with cross-type loading support. + + Args: + path: The file path to load from + target_grid_class: Optional Grid class to load into. If None, uses default Grid. + + Returns: + Grid: The deserialized Grid object of the specified target class + """ + with open(path, "rb") as f: + data = f.read() + + # Extract extensions and deserialize core data + input_data, extensions = _extract_msgpack_data(data) + + # Create grid and restore extensions + grid = _create_grid_from_input_data(input_data, target_grid_class) + _restore_extensions_data(grid, extensions) + + return grid diff --git a/tests/unit/utils/test_serialization.py b/tests/unit/utils/test_serialization.py new file mode 100644 index 0000000..09f988f --- /dev/null +++ b/tests/unit/utils/test_serialization.py @@ -0,0 +1,548 @@ +# SPDX-FileCopyrightText: Contributors to the Power Grid Model project +# +# SPDX-License-Identifier: MPL-2.0 + +"""Comprehensive unit tests for Grid serialization with power-grid-model compatibility.""" + +import json +from dataclasses import dataclass +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import patch + +import numpy as np +import pytest +from numpy.typing import NDArray + +from power_grid_model_ds import Grid +from power_grid_model_ds._core.load_flow import PowerGridModelInterface +from power_grid_model_ds._core.model.arrays.base.array import FancyArray +from power_grid_model_ds._core.utils.serialization import ( + _extract_extensions_data, + _get_serialization_path, + _restore_extensions_data, + load_grid_from_json, + load_grid_from_msgpack, + save_grid_to_json, + save_grid_to_msgpack, +) +from power_grid_model_ds.arrays import LineArray +from power_grid_model_ds.arrays import NodeArray as BaseNodeArray + + +class ExtendedNodeArray(BaseNodeArray): + """Test array with extended columns""" + + _defaults = {"u": 0.0, "analysis_flag": 0} + u: NDArray[np.float64] + analysis_flag: NDArray[np.int32] + + +class ExtendedLineArray(LineArray): + """Test array with extended columns""" + + _defaults = {"i_from": 0.0, "loading_factor": 0.0} + i_from: NDArray[np.float64] + loading_factor: NDArray[np.float64] + + +@dataclass +class ExtendedGrid(Grid): + """Test grid with extended arrays""" + + node: ExtendedNodeArray + line: ExtendedLineArray + + +@pytest.fixture +def temp_dir(): + """Temporary directory fixture""" + with TemporaryDirectory() as tmp_dir: + yield Path(tmp_dir) + + +@pytest.fixture +def basic_grid(): + """Basic grid fixture""" + return Grid.from_txt("1 2", "2 3", "S10 1") + + +@pytest.fixture +def extended_grid(): + """Extended grid fixture with additional columns""" + grid = ExtendedGrid.empty() + nodes = ExtendedNodeArray( + id=[1, 2, 3], u_rated=[10500, 10500, 10500], u=[10450, 10400, 10350], analysis_flag=[1, 0, 1] + ) + lines = ExtendedLineArray( + id=[10, 11], + from_node=[1, 2], + to_node=[2, 3], + from_status=[1, 1], + to_status=[1, 1], + r1=[0.1, 0.15], + x1=[0.2, 0.25], + c1=[1e-6, 1.2e-6], + tan1=[0.0, 0.0], + i_n=[400, 350], + i_from=[150.5, 120.3], + loading_factor=[0.75, 0.68], + ) + grid.append(nodes) + grid.append(lines) + return grid + + +class TestSerializationFormats: + """Test serialization across different formats and configurations""" + + @pytest.mark.parametrize( + "format_type,preserve_ext", [("json", True), ("json", False), ("msgpack", True), ("msgpack", False)] + ) + def test_basic_serialization_roundtrip( + self, basic_grid: Grid, temp_dir: Path, format_type: str, preserve_ext: bool + ): + """Test basic serialization roundtrip for all formats""" + ext = "json" if format_type == "json" else "msgpack" + path = temp_dir / f"test.{ext}" + + # Save + if format_type == "json": + result_path = save_grid_to_json(basic_grid, path, preserve_extensions=preserve_ext) + else: + result_path = save_grid_to_msgpack(basic_grid, path, preserve_extensions=preserve_ext) + + assert result_path.exists() + + # Load and verify + if format_type == "json": + loaded_grid = load_grid_from_json(path, target_grid_class=Grid) + else: + loaded_grid = load_grid_from_msgpack(path, target_grid_class=Grid) + assert loaded_grid.node.size == basic_grid.node.size + assert loaded_grid.line.size == basic_grid.line.size + assert list(loaded_grid.node.id) == list(basic_grid.node.id) + + @pytest.mark.parametrize("format_type", ["json", "msgpack"]) + def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, temp_dir: Path, format_type: str): + """Test extended serialization preserving custom data""" + ext = "json" if format_type == "json" else "msgpack" + path = temp_dir / f"extended.{ext}" + + # Save with extensions + if format_type == "json": + save_grid_to_json(extended_grid, path, preserve_extensions=True) + loaded_grid = load_grid_from_json(path, target_grid_class=ExtendedGrid) + else: + save_grid_to_msgpack(extended_grid, path, preserve_extensions=True) + loaded_grid = load_grid_from_msgpack(path, target_grid_class=ExtendedGrid) + + # Verify core data + assert loaded_grid.node.size == extended_grid.node.size + assert loaded_grid.line.size == extended_grid.line.size + + # Verify extended data + np.testing.assert_array_equal(loaded_grid.node.u, extended_grid.node.u) + np.testing.assert_array_equal(loaded_grid.line.i_from, extended_grid.line.i_from) + + def test_format_characteristics(self, basic_grid: Grid, temp_dir: Path): + """Test format-specific characteristics""" + json_path = temp_dir / "test.json" + msgpack_path = temp_dir / "test.msgpack" + + save_grid_to_json(basic_grid, json_path) + save_grid_to_msgpack(basic_grid, msgpack_path) + + # JSON should be human-readable + with open(json_path) as f: + json_data = json.load(f) + assert "version" in json_data + assert "data" in json_data + + # MessagePack should be more compact + json_size = json_path.stat().st_size + msgpack_size = msgpack_path.stat().st_size + assert msgpack_size < json_size # MessagePack is typically more compact + + +class TestCrossTypeCompatibility: + """Test cross-type loading and compatibility""" + + @pytest.mark.parametrize("format_type", ["json", "msgpack"]) + def test_basic_to_extended_loading(self, basic_grid: Grid, temp_dir: Path, format_type: str): + """Test loading basic grid into extended type""" + ext = "json" if format_type == "json" else "msgpack" + path = temp_dir / f"basic.{ext}" + + # Save basic grid + if format_type == "json": + save_grid_to_json(basic_grid, path) + loaded_grid = load_grid_from_json(path, target_grid_class=ExtendedGrid) + else: + save_grid_to_msgpack(basic_grid, path) + loaded_grid = load_grid_from_msgpack(path, target_grid_class=ExtendedGrid) + + # Core data should transfer + assert loaded_grid.node.size == basic_grid.node.size + assert loaded_grid.line.size == basic_grid.line.size + + @pytest.mark.parametrize("format_type", ["json", "msgpack"]) + def test_extended_to_basic_loading(self, extended_grid: ExtendedGrid, temp_dir: Path, format_type: str): + """Test loading extended grid into basic type""" + ext = "json" if format_type == "json" else "msgpack" + path = temp_dir / f"extended.{ext}" + + # Save extended grid + if format_type == "json": + save_grid_to_json(extended_grid, path, preserve_extensions=True) + loaded_grid = load_grid_from_json(path, target_grid_class=Grid) + else: + save_grid_to_msgpack(extended_grid, path, preserve_extensions=True) + loaded_grid = load_grid_from_msgpack(path, target_grid_class=Grid) + + # Core data should transfer + assert loaded_grid.node.size == extended_grid.node.size + assert loaded_grid.line.size == extended_grid.line.size + + +class TestExtensionHandling: + """Test extension data handling and edge cases""" + + def test_extension_extraction_and_restoration(self, extended_grid: ExtendedGrid): + """Test extension data extraction and restoration""" + # Extract extensions + extensions = _extract_extensions_data(extended_grid) + + # Should have both types of extensions + assert "extended_columns" in extensions + assert "custom_arrays" in extensions + + # Should contain extended node data + assert "node" in extensions["extended_columns"] + assert "u" in extensions["extended_columns"]["node"] + + # Create new grid and restore + new_grid = ExtendedGrid.empty() + # Add basic structure first + nodes = ExtendedNodeArray(id=[1, 2, 3], u_rated=[10500, 10500, 10500]) + new_grid.append(nodes) + + # Restore extensions + _restore_extensions_data(new_grid, extensions) + + # Extended data should be restored + np.testing.assert_array_equal(new_grid.node.u, extended_grid.node.u) + + def test_missing_extension_keys(self): + """Test graceful handling of missing extension keys""" + basic_grid = Grid.empty() + + # Test various malformed extension data + test_cases = [ + {}, # Empty + {"extended_columns": {}}, # Missing custom_arrays + {"custom_arrays": {}}, # Missing extended_columns + {"extended_columns": {"test": "value"}}, # Invalid structure + ] + + for extensions in test_cases: + # Should not raise + _restore_extensions_data(basic_grid, extensions) + + def test_custom_array_serialization_roundtrip(self, temp_dir: Path): + """Test serialization and loading of grids with custom arrays""" + + # Create a custom array type that properly extends FancyArray + class CustomMetadataArray(FancyArray): + """Custom metadata array for testing""" + + _defaults = {"metadata_value": 0.0, "category": 0} + + id: NDArray[np.int32] + metadata_value: NDArray[np.float64] + category: NDArray[np.int32] + + # Create a grid with custom arrays + @dataclass + class GridWithCustomArray(Grid): + custom_metadata: CustomMetadataArray + + # Create test grid with custom data + grid = GridWithCustomArray.empty() + + # Add some basic grid data + nodes = grid.node.__class__(id=[1, 2], u_rated=[10000, 10000]) + grid.append(nodes) + + # Add custom metadata + custom_data = CustomMetadataArray(id=[100, 200, 300], metadata_value=[1.5, 2.5, 3.5], category=[1, 2, 1]) + grid.custom_metadata = custom_data + + # Test JSON serialization + json_path = temp_dir / "custom_array.json" + save_grid_to_json(grid, json_path, preserve_extensions=True) + + # Load back and verify + loaded_grid = load_grid_from_json(json_path, target_grid_class=GridWithCustomArray) + + # Verify core data + assert loaded_grid.node.size == 2 + np.testing.assert_array_equal(loaded_grid.node.id, [1, 2]) + + # Verify custom array was preserved + assert hasattr(loaded_grid, "custom_metadata") + assert loaded_grid.custom_metadata.size == 3 + np.testing.assert_array_equal(loaded_grid.custom_metadata.id, [100, 200, 300]) + np.testing.assert_array_almost_equal(loaded_grid.custom_metadata.metadata_value, [1.5, 2.5, 3.5]) + np.testing.assert_array_equal(loaded_grid.custom_metadata.category, [1, 2, 1]) + + # Test MessagePack serialization + msgpack_path = temp_dir / "custom_array.msgpack" + save_grid_to_msgpack(grid, msgpack_path, preserve_extensions=True) + + # Load back and verify + loaded_grid_mp = load_grid_from_msgpack(msgpack_path, target_grid_class=GridWithCustomArray) + + # Verify core data + assert loaded_grid_mp.node.size == 2 + np.testing.assert_array_equal(loaded_grid_mp.node.id, [1, 2]) + + # Verify custom array was preserved + assert hasattr(loaded_grid_mp, "custom_metadata") + assert loaded_grid_mp.custom_metadata.size == 3 + np.testing.assert_array_equal(loaded_grid_mp.custom_metadata.id, [100, 200, 300]) + np.testing.assert_array_almost_equal(loaded_grid_mp.custom_metadata.metadata_value, [1.5, 2.5, 3.5]) + np.testing.assert_array_equal(loaded_grid_mp.custom_metadata.category, [1, 2, 1]) + + +class TestErrorHandling: + """Test error handling and edge cases""" + + @pytest.mark.parametrize("format_type", ["json", "msgpack"]) + def test_file_not_found(self, format_type: str): + """Test handling of missing files""" + nonexistent_path = Path("nonexistent.") / format_type + + with pytest.raises(FileNotFoundError): + if format_type == "json": + load_grid_from_json(nonexistent_path) + else: + load_grid_from_msgpack(nonexistent_path) + + def test_corrupted_json_content(self, temp_dir: Path): + """Test handling of corrupted JSON files""" + json_path = temp_dir / "corrupted.json" + + with open(json_path, "w") as f: + f.write("{ invalid json content") + + with pytest.raises(json.JSONDecodeError): + load_grid_from_json(json_path) + + def test_corrupted_msgpack_data(self, temp_dir: Path): + """Test handling of corrupted MessagePack files""" + msgpack_path = temp_dir / "corrupted.msgpack" + + with open(msgpack_path, "wb") as f: + f.write(b"invalid msgpack data") + + with pytest.raises(Exception): # Could be various msgpack/power-grid-model exceptions + load_grid_from_msgpack(msgpack_path) + + def test_invalid_extension_data_recovery(self, temp_dir: Path): + """Test recovery from invalid extension data""" + # Create valid extended grid + extended_grid = ExtendedGrid.empty() + nodes = ExtendedNodeArray(id=[1, 2], u_rated=[10000, 10000], u=[9950, 9900]) + extended_grid.append(nodes) + + json_path = temp_dir / "test_recovery.json" + save_grid_to_json(extended_grid, json_path, preserve_extensions=True) + + # Corrupt extension data + with open(json_path, "r") as f: + data = json.load(f) + + # Add invalid extension data + if "pgm_ds_extensions" in data: + data["pgm_ds_extensions"]["extended_columns"]["node"]["u"] = [1, 2, 3, 4, 5] # Wrong size + data["pgm_ds_extensions"]["custom_arrays"]["fake"] = {"dtype": "invalid_dtype", "data": [[1, 2, 3]]} + + with open(json_path, "w") as f: + json.dump(data, f) + + # Should load core data despite extension errors + loaded_grid = load_grid_from_json(json_path, target_grid_class=Grid) + assert loaded_grid.node.size == 2 + + def test_io_permission_errors(self, temp_dir: Path): + """Test I/O permission error handling""" + basic_grid = Grid.empty() + + # Create read-only directory + readonly_dir = temp_dir / "readonly" + readonly_dir.mkdir() + readonly_dir.chmod(0o444) + + readonly_path = readonly_dir / "test.json" + + try: + with pytest.raises((PermissionError, OSError)): + save_grid_to_json(basic_grid, readonly_path) + finally: + readonly_dir.chmod(0o755) # Cleanup + + +class TestUtilityFunctions: + """Test utility functions and path handling""" + + @pytest.mark.parametrize( + "input_path,format_type,expected", + [ + ("test.json", "auto", "test.json"), + ("test.msgpack", "auto", "test.msgpack"), + ("test.mp", "auto", "test.mp"), + ("test.xyz", "auto", "test.json"), # Unknown defaults to JSON + ("test.xyz", "json", "test.json"), + ("test.xyz", "msgpack", "test.msgpack"), + ], + ) + def test_serialization_path_handling(self, input_path: str, format_type: str, expected: str): + """Test path handling and format detection""" + result = _get_serialization_path(Path(input_path), format_type) + assert result == Path(expected) + + +class TestPowerGridModelIntegration: + """Test integration with PowerGridModelInterface""" + + def test_interface_roundtrip(self, basic_grid: Grid, temp_dir: Path): + """Test roundtrip through PowerGridModelInterface""" + json_path = temp_dir / "interface_test.json" + + # Save and load through interface + save_grid_to_json(basic_grid, json_path) + + # Should work with interface directly + interface = PowerGridModelInterface(grid=basic_grid) + input_data = interface.create_input_from_grid() + + # Create new grid from input data + new_interface = PowerGridModelInterface(input_data=input_data) + new_grid = new_interface.create_grid_from_input_data() + + assert new_grid.node.size == basic_grid.node.size + + def test_interface_error_propagation(self, temp_dir: Path): + """Test error propagation from PowerGridModelInterface""" + # This is harder to test directly, but we can verify errors are not swallowed + basic_grid = Grid.empty() + json_path = temp_dir / "error_test.json" + + # Create conditions that might cause interface errors + with patch.object(PowerGridModelInterface, "create_input_from_grid", side_effect=Exception("PGM Error")): + with pytest.raises(Exception, match="PGM Error"): + save_grid_to_json(basic_grid, json_path) + + +class TestSpecialCases: + """Test special cases and edge scenarios""" + + def test_empty_grid_handling(self, temp_dir: Path): + """Test serialization of empty grids""" + empty_grid = Grid.empty() + + json_path = temp_dir / "empty.json" + msgpack_path = temp_dir / "empty.msgpack" + + # Should handle empty grids + save_grid_to_json(empty_grid, json_path) + save_grid_to_msgpack(empty_grid, msgpack_path) + + # Should load back as empty + loaded_json = load_grid_from_json(json_path, target_grid_class=Grid) + loaded_msgpack = load_grid_from_msgpack(msgpack_path, target_grid_class=Grid) + + assert loaded_json.node.size == 0 + assert loaded_msgpack.node.size == 0 + + def test_extreme_values_handling(self, temp_dir: Path): + """Test handling of extreme numeric values""" + extended_grid = ExtendedGrid.empty() + + # Add nodes with extreme values + nodes = ExtendedNodeArray( + id=[1, 2], + u_rated=[10000, 10000], + u=[np.inf, -np.inf], # Extreme values + ) + extended_grid.append(nodes) + + json_path = temp_dir / "extreme.json" + + # Should handle extreme values (JSON supports inf) + save_grid_to_json(extended_grid, json_path, preserve_extensions=True) + loaded_grid = load_grid_from_json(json_path, target_grid_class=Grid) + assert loaded_grid.node.size == 2 + + def test_without_extensions_flag(self, extended_grid: ExtendedGrid, temp_dir: Path): + """Test serialization without extensions flag""" + json_path = temp_dir / "no_ext.json" + msgpack_path = temp_dir / "no_ext.msgpack" + + # Save without extensions + save_grid_to_json(extended_grid, json_path, preserve_extensions=False) + save_grid_to_msgpack(extended_grid, msgpack_path, preserve_extensions=False) + + # Should load core data only + loaded_json = load_grid_from_json(json_path, target_grid_class=Grid) + loaded_msgpack = load_grid_from_msgpack(msgpack_path, target_grid_class=Grid) + + assert loaded_json.node.size == extended_grid.node.size + assert loaded_msgpack.node.size == extended_grid.node.size + + def test_directory_creation_during_save(self, basic_grid: Grid, temp_dir: Path): + """Test automatic directory creation during save operations""" + # Test nested directory creation + nested_json_path = temp_dir / "nested" / "deep" / "test.json" + nested_msgpack_path = temp_dir / "nested" / "deep" / "test.msgpack" + + # Should create directories automatically + save_grid_to_json(basic_grid, nested_json_path) + save_grid_to_msgpack(basic_grid, nested_msgpack_path) + + assert nested_json_path.exists() + assert nested_msgpack_path.exists() + + # Should be able to load back + loaded_json = load_grid_from_json(nested_json_path, target_grid_class=Grid) + loaded_msgpack = load_grid_from_msgpack(nested_msgpack_path, target_grid_class=Grid) + + assert loaded_json.node.size == basic_grid.node.size + assert loaded_msgpack.node.size == basic_grid.node.size + + def test_custom_array_extraction_edge_cases(self, temp_dir: Path): + """Test edge cases in custom array extraction""" + # Test with grid that has complex custom arrays that might cause extraction issues + extended_grid = ExtendedGrid.empty() + + # Add data that might cause issues during extraction + nodes = ExtendedNodeArray( + id=[1, 2], + u_rated=[10000, 10000], + u=[float("nan"), float("inf")], # Edge case values + ) + extended_grid.append(nodes) + + # Should handle edge case values gracefully + extensions = _extract_extensions_data(extended_grid) + assert "extended_columns" in extensions + assert "custom_arrays" in extensions + + # Test saving and loading with these edge cases + json_path = temp_dir / "edge_cases.json" + save_grid_to_json(extended_grid, json_path, preserve_extensions=True) + + # Should load without issues + loaded_grid = load_grid_from_json(json_path, target_grid_class=Grid) + assert loaded_grid.node.size == 2 From 1491e1401fe928622cf736b869432e7299a7569e Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Mon, 13 Oct 2025 19:09:08 +0200 Subject: [PATCH 02/19] test: cleanup tests Signed-off-by: jaapschoutenalliander --- tests/unit/utils/test_serialization.py | 237 +++---------------------- 1 file changed, 26 insertions(+), 211 deletions(-) diff --git a/tests/unit/utils/test_serialization.py b/tests/unit/utils/test_serialization.py index 09f988f..5e0ad61 100644 --- a/tests/unit/utils/test_serialization.py +++ b/tests/unit/utils/test_serialization.py @@ -8,14 +8,12 @@ from dataclasses import dataclass from pathlib import Path from tempfile import TemporaryDirectory -from unittest.mock import patch import numpy as np import pytest from numpy.typing import NDArray from power_grid_model_ds import Grid -from power_grid_model_ds._core.load_flow import PowerGridModelInterface from power_grid_model_ds._core.model.arrays.base.array import FancyArray from power_grid_model_ds._core.utils.serialization import ( _extract_extensions_data, @@ -145,25 +143,6 @@ def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, tem np.testing.assert_array_equal(loaded_grid.node.u, extended_grid.node.u) np.testing.assert_array_equal(loaded_grid.line.i_from, extended_grid.line.i_from) - def test_format_characteristics(self, basic_grid: Grid, temp_dir: Path): - """Test format-specific characteristics""" - json_path = temp_dir / "test.json" - msgpack_path = temp_dir / "test.msgpack" - - save_grid_to_json(basic_grid, json_path) - save_grid_to_msgpack(basic_grid, msgpack_path) - - # JSON should be human-readable - with open(json_path) as f: - json_data = json.load(f) - assert "version" in json_data - assert "data" in json_data - - # MessagePack should be more compact - json_size = json_path.stat().st_size - msgpack_size = msgpack_path.stat().st_size - assert msgpack_size < json_size # MessagePack is typically more compact - class TestCrossTypeCompatibility: """Test cross-type loading and compatibility""" @@ -208,31 +187,6 @@ def test_extended_to_basic_loading(self, extended_grid: ExtendedGrid, temp_dir: class TestExtensionHandling: """Test extension data handling and edge cases""" - def test_extension_extraction_and_restoration(self, extended_grid: ExtendedGrid): - """Test extension data extraction and restoration""" - # Extract extensions - extensions = _extract_extensions_data(extended_grid) - - # Should have both types of extensions - assert "extended_columns" in extensions - assert "custom_arrays" in extensions - - # Should contain extended node data - assert "node" in extensions["extended_columns"] - assert "u" in extensions["extended_columns"]["node"] - - # Create new grid and restore - new_grid = ExtendedGrid.empty() - # Add basic structure first - nodes = ExtendedNodeArray(id=[1, 2, 3], u_rated=[10500, 10500, 10500]) - new_grid.append(nodes) - - # Restore extensions - _restore_extensions_data(new_grid, extensions) - - # Extended data should be restored - np.testing.assert_array_equal(new_grid.node.u, extended_grid.node.u) - def test_missing_extension_keys(self): """Test graceful handling of missing extension keys""" basic_grid = Grid.empty() @@ -315,84 +269,6 @@ class GridWithCustomArray(Grid): np.testing.assert_array_equal(loaded_grid_mp.custom_metadata.category, [1, 2, 1]) -class TestErrorHandling: - """Test error handling and edge cases""" - - @pytest.mark.parametrize("format_type", ["json", "msgpack"]) - def test_file_not_found(self, format_type: str): - """Test handling of missing files""" - nonexistent_path = Path("nonexistent.") / format_type - - with pytest.raises(FileNotFoundError): - if format_type == "json": - load_grid_from_json(nonexistent_path) - else: - load_grid_from_msgpack(nonexistent_path) - - def test_corrupted_json_content(self, temp_dir: Path): - """Test handling of corrupted JSON files""" - json_path = temp_dir / "corrupted.json" - - with open(json_path, "w") as f: - f.write("{ invalid json content") - - with pytest.raises(json.JSONDecodeError): - load_grid_from_json(json_path) - - def test_corrupted_msgpack_data(self, temp_dir: Path): - """Test handling of corrupted MessagePack files""" - msgpack_path = temp_dir / "corrupted.msgpack" - - with open(msgpack_path, "wb") as f: - f.write(b"invalid msgpack data") - - with pytest.raises(Exception): # Could be various msgpack/power-grid-model exceptions - load_grid_from_msgpack(msgpack_path) - - def test_invalid_extension_data_recovery(self, temp_dir: Path): - """Test recovery from invalid extension data""" - # Create valid extended grid - extended_grid = ExtendedGrid.empty() - nodes = ExtendedNodeArray(id=[1, 2], u_rated=[10000, 10000], u=[9950, 9900]) - extended_grid.append(nodes) - - json_path = temp_dir / "test_recovery.json" - save_grid_to_json(extended_grid, json_path, preserve_extensions=True) - - # Corrupt extension data - with open(json_path, "r") as f: - data = json.load(f) - - # Add invalid extension data - if "pgm_ds_extensions" in data: - data["pgm_ds_extensions"]["extended_columns"]["node"]["u"] = [1, 2, 3, 4, 5] # Wrong size - data["pgm_ds_extensions"]["custom_arrays"]["fake"] = {"dtype": "invalid_dtype", "data": [[1, 2, 3]]} - - with open(json_path, "w") as f: - json.dump(data, f) - - # Should load core data despite extension errors - loaded_grid = load_grid_from_json(json_path, target_grid_class=Grid) - assert loaded_grid.node.size == 2 - - def test_io_permission_errors(self, temp_dir: Path): - """Test I/O permission error handling""" - basic_grid = Grid.empty() - - # Create read-only directory - readonly_dir = temp_dir / "readonly" - readonly_dir.mkdir() - readonly_dir.chmod(0o444) - - readonly_path = readonly_dir / "test.json" - - try: - with pytest.raises((PermissionError, OSError)): - save_grid_to_json(basic_grid, readonly_path) - finally: - readonly_dir.chmod(0o755) # Cleanup - - class TestUtilityFunctions: """Test utility functions and path handling""" @@ -413,38 +289,6 @@ def test_serialization_path_handling(self, input_path: str, format_type: str, ex assert result == Path(expected) -class TestPowerGridModelIntegration: - """Test integration with PowerGridModelInterface""" - - def test_interface_roundtrip(self, basic_grid: Grid, temp_dir: Path): - """Test roundtrip through PowerGridModelInterface""" - json_path = temp_dir / "interface_test.json" - - # Save and load through interface - save_grid_to_json(basic_grid, json_path) - - # Should work with interface directly - interface = PowerGridModelInterface(grid=basic_grid) - input_data = interface.create_input_from_grid() - - # Create new grid from input data - new_interface = PowerGridModelInterface(input_data=input_data) - new_grid = new_interface.create_grid_from_input_data() - - assert new_grid.node.size == basic_grid.node.size - - def test_interface_error_propagation(self, temp_dir: Path): - """Test error propagation from PowerGridModelInterface""" - # This is harder to test directly, but we can verify errors are not swallowed - basic_grid = Grid.empty() - json_path = temp_dir / "error_test.json" - - # Create conditions that might cause interface errors - with patch.object(PowerGridModelInterface, "create_input_from_grid", side_effect=Exception("PGM Error")): - with pytest.raises(Exception, match="PGM Error"): - save_grid_to_json(basic_grid, json_path) - - class TestSpecialCases: """Test special cases and edge scenarios""" @@ -466,61 +310,6 @@ def test_empty_grid_handling(self, temp_dir: Path): assert loaded_json.node.size == 0 assert loaded_msgpack.node.size == 0 - def test_extreme_values_handling(self, temp_dir: Path): - """Test handling of extreme numeric values""" - extended_grid = ExtendedGrid.empty() - - # Add nodes with extreme values - nodes = ExtendedNodeArray( - id=[1, 2], - u_rated=[10000, 10000], - u=[np.inf, -np.inf], # Extreme values - ) - extended_grid.append(nodes) - - json_path = temp_dir / "extreme.json" - - # Should handle extreme values (JSON supports inf) - save_grid_to_json(extended_grid, json_path, preserve_extensions=True) - loaded_grid = load_grid_from_json(json_path, target_grid_class=Grid) - assert loaded_grid.node.size == 2 - - def test_without_extensions_flag(self, extended_grid: ExtendedGrid, temp_dir: Path): - """Test serialization without extensions flag""" - json_path = temp_dir / "no_ext.json" - msgpack_path = temp_dir / "no_ext.msgpack" - - # Save without extensions - save_grid_to_json(extended_grid, json_path, preserve_extensions=False) - save_grid_to_msgpack(extended_grid, msgpack_path, preserve_extensions=False) - - # Should load core data only - loaded_json = load_grid_from_json(json_path, target_grid_class=Grid) - loaded_msgpack = load_grid_from_msgpack(msgpack_path, target_grid_class=Grid) - - assert loaded_json.node.size == extended_grid.node.size - assert loaded_msgpack.node.size == extended_grid.node.size - - def test_directory_creation_during_save(self, basic_grid: Grid, temp_dir: Path): - """Test automatic directory creation during save operations""" - # Test nested directory creation - nested_json_path = temp_dir / "nested" / "deep" / "test.json" - nested_msgpack_path = temp_dir / "nested" / "deep" / "test.msgpack" - - # Should create directories automatically - save_grid_to_json(basic_grid, nested_json_path) - save_grid_to_msgpack(basic_grid, nested_msgpack_path) - - assert nested_json_path.exists() - assert nested_msgpack_path.exists() - - # Should be able to load back - loaded_json = load_grid_from_json(nested_json_path, target_grid_class=Grid) - loaded_msgpack = load_grid_from_msgpack(nested_msgpack_path, target_grid_class=Grid) - - assert loaded_json.node.size == basic_grid.node.size - assert loaded_msgpack.node.size == basic_grid.node.size - def test_custom_array_extraction_edge_cases(self, temp_dir: Path): """Test edge cases in custom array extraction""" # Test with grid that has complex custom arrays that might cause extraction issues @@ -546,3 +335,29 @@ def test_custom_array_extraction_edge_cases(self, temp_dir: Path): # Should load without issues loaded_grid = load_grid_from_json(json_path, target_grid_class=Grid) assert loaded_grid.node.size == 2 + + def test_invalid_extension_data_recovery(self, temp_dir: Path): + """Test recovery from invalid extension data""" + # Create valid extended grid + extended_grid = ExtendedGrid.empty() + nodes = ExtendedNodeArray(id=[1, 2], u_rated=[10000, 10000], u=[9950, 9900]) + extended_grid.append(nodes) + + json_path = temp_dir / "test_recovery.json" + save_grid_to_json(extended_grid, json_path, preserve_extensions=True) + + # Corrupt extension data + with open(json_path, "r") as f: + data = json.load(f) + + # Add invalid extension data + if "pgm_ds_extensions" in data: + data["pgm_ds_extensions"]["extended_columns"]["node"]["u"] = [1, 2, 3, 4, 5] # Wrong size + data["pgm_ds_extensions"]["custom_arrays"]["fake"] = {"dtype": "invalid_dtype", "data": [[1, 2, 3]]} + + with open(json_path, "w") as f: + json.dump(data, f) + + # Should load core data despite extension errors + loaded_grid = load_grid_from_json(json_path, target_grid_class=Grid) + assert loaded_grid.node.size == 2 From 6e11a3de74afa1eab29492538092c5cbda2c4330 Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Mon, 13 Oct 2025 19:11:23 +0200 Subject: [PATCH 03/19] chore: requirements add msgpack Signed-off-by: jaapschoutenalliander --- pyproject.toml | 1 + uv.lock | 40 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 04b1aaf..110dab2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,7 @@ dependencies = [ "power-grid-model>=1.7", "rustworkx>= 0.15.1", "numpy>=2.0", + "msgpack>=1.1.2", ] dynamic = ["version"] diff --git a/uv.lock b/uv.lock index 27a131b..5346f13 100644 --- a/uv.lock +++ b/uv.lock @@ -763,6 +763,41 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] +[[package]] +name = "msgpack" +version = "1.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4d/f2/bfb55a6236ed8725a96b0aa3acbd0ec17588e6a2c3b62a93eb513ed8783f/msgpack-1.1.2.tar.gz", hash = "sha256:3b60763c1373dd60f398488069bcdc703cd08a711477b5d480eecc9f9626f47e", size = 173581, upload-time = "2025-10-08T09:15:56.596Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/97/560d11202bcd537abca693fd85d81cebe2107ba17301de42b01ac1677b69/msgpack-1.1.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:2e86a607e558d22985d856948c12a3fa7b42efad264dca8a3ebbcfa2735d786c", size = 82271, upload-time = "2025-10-08T09:14:49.967Z" }, + { url = "https://files.pythonhosted.org/packages/83/04/28a41024ccbd67467380b6fb440ae916c1e4f25e2cd4c63abe6835ac566e/msgpack-1.1.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:283ae72fc89da59aa004ba147e8fc2f766647b1251500182fac0350d8af299c0", size = 84914, upload-time = "2025-10-08T09:14:50.958Z" }, + { url = "https://files.pythonhosted.org/packages/71/46/b817349db6886d79e57a966346cf0902a426375aadc1e8e7a86a75e22f19/msgpack-1.1.2-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:61c8aa3bd513d87c72ed0b37b53dd5c5a0f58f2ff9f26e1555d3bd7948fb7296", size = 416962, upload-time = "2025-10-08T09:14:51.997Z" }, + { url = "https://files.pythonhosted.org/packages/da/e0/6cc2e852837cd6086fe7d8406af4294e66827a60a4cf60b86575a4a65ca8/msgpack-1.1.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:454e29e186285d2ebe65be34629fa0e8605202c60fbc7c4c650ccd41870896ef", size = 426183, upload-time = "2025-10-08T09:14:53.477Z" }, + { url = "https://files.pythonhosted.org/packages/25/98/6a19f030b3d2ea906696cedd1eb251708e50a5891d0978b012cb6107234c/msgpack-1.1.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:7bc8813f88417599564fafa59fd6f95be417179f76b40325b500b3c98409757c", size = 411454, upload-time = "2025-10-08T09:14:54.648Z" }, + { url = "https://files.pythonhosted.org/packages/b7/cd/9098fcb6adb32187a70b7ecaabf6339da50553351558f37600e53a4a2a23/msgpack-1.1.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:bafca952dc13907bdfdedfc6a5f579bf4f292bdd506fadb38389afa3ac5b208e", size = 422341, upload-time = "2025-10-08T09:14:56.328Z" }, + { url = "https://files.pythonhosted.org/packages/e6/ae/270cecbcf36c1dc85ec086b33a51a4d7d08fc4f404bdbc15b582255d05ff/msgpack-1.1.2-cp311-cp311-win32.whl", hash = "sha256:602b6740e95ffc55bfb078172d279de3773d7b7db1f703b2f1323566b878b90e", size = 64747, upload-time = "2025-10-08T09:14:57.882Z" }, + { url = "https://files.pythonhosted.org/packages/2a/79/309d0e637f6f37e83c711f547308b91af02b72d2326ddd860b966080ef29/msgpack-1.1.2-cp311-cp311-win_amd64.whl", hash = "sha256:d198d275222dc54244bf3327eb8cbe00307d220241d9cec4d306d49a44e85f68", size = 71633, upload-time = "2025-10-08T09:14:59.177Z" }, + { url = "https://files.pythonhosted.org/packages/73/4d/7c4e2b3d9b1106cd0aa6cb56cc57c6267f59fa8bfab7d91df5adc802c847/msgpack-1.1.2-cp311-cp311-win_arm64.whl", hash = "sha256:86f8136dfa5c116365a8a651a7d7484b65b13339731dd6faebb9a0242151c406", size = 64755, upload-time = "2025-10-08T09:15:00.48Z" }, + { url = "https://files.pythonhosted.org/packages/ad/bd/8b0d01c756203fbab65d265859749860682ccd2a59594609aeec3a144efa/msgpack-1.1.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:70a0dff9d1f8da25179ffcf880e10cf1aad55fdb63cd59c9a49a1b82290062aa", size = 81939, upload-time = "2025-10-08T09:15:01.472Z" }, + { url = "https://files.pythonhosted.org/packages/34/68/ba4f155f793a74c1483d4bdef136e1023f7bcba557f0db4ef3db3c665cf1/msgpack-1.1.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:446abdd8b94b55c800ac34b102dffd2f6aa0ce643c55dfc017ad89347db3dbdb", size = 85064, upload-time = "2025-10-08T09:15:03.764Z" }, + { url = "https://files.pythonhosted.org/packages/f2/60/a064b0345fc36c4c3d2c743c82d9100c40388d77f0b48b2f04d6041dbec1/msgpack-1.1.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c63eea553c69ab05b6747901b97d620bb2a690633c77f23feb0c6a947a8a7b8f", size = 417131, upload-time = "2025-10-08T09:15:05.136Z" }, + { url = "https://files.pythonhosted.org/packages/65/92/a5100f7185a800a5d29f8d14041f61475b9de465ffcc0f3b9fba606e4505/msgpack-1.1.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:372839311ccf6bdaf39b00b61288e0557916c3729529b301c52c2d88842add42", size = 427556, upload-time = "2025-10-08T09:15:06.837Z" }, + { url = "https://files.pythonhosted.org/packages/f5/87/ffe21d1bf7d9991354ad93949286f643b2bb6ddbeab66373922b44c3b8cc/msgpack-1.1.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2929af52106ca73fcb28576218476ffbb531a036c2adbcf54a3664de124303e9", size = 404920, upload-time = "2025-10-08T09:15:08.179Z" }, + { url = "https://files.pythonhosted.org/packages/ff/41/8543ed2b8604f7c0d89ce066f42007faac1eaa7d79a81555f206a5cdb889/msgpack-1.1.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:be52a8fc79e45b0364210eef5234a7cf8d330836d0a64dfbb878efa903d84620", size = 415013, upload-time = "2025-10-08T09:15:09.83Z" }, + { url = "https://files.pythonhosted.org/packages/41/0d/2ddfaa8b7e1cee6c490d46cb0a39742b19e2481600a7a0e96537e9c22f43/msgpack-1.1.2-cp312-cp312-win32.whl", hash = "sha256:1fff3d825d7859ac888b0fbda39a42d59193543920eda9d9bea44d958a878029", size = 65096, upload-time = "2025-10-08T09:15:11.11Z" }, + { url = "https://files.pythonhosted.org/packages/8c/ec/d431eb7941fb55a31dd6ca3404d41fbb52d99172df2e7707754488390910/msgpack-1.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:1de460f0403172cff81169a30b9a92b260cb809c4cb7e2fc79ae8d0510c78b6b", size = 72708, upload-time = "2025-10-08T09:15:12.554Z" }, + { url = "https://files.pythonhosted.org/packages/c5/31/5b1a1f70eb0e87d1678e9624908f86317787b536060641d6798e3cf70ace/msgpack-1.1.2-cp312-cp312-win_arm64.whl", hash = "sha256:be5980f3ee0e6bd44f3a9e9dea01054f175b50c3e6cdb692bc9424c0bbb8bf69", size = 64119, upload-time = "2025-10-08T09:15:13.589Z" }, + { url = "https://files.pythonhosted.org/packages/6b/31/b46518ecc604d7edf3a4f94cb3bf021fc62aa301f0cb849936968164ef23/msgpack-1.1.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:4efd7b5979ccb539c221a4c4e16aac1a533efc97f3b759bb5a5ac9f6d10383bf", size = 81212, upload-time = "2025-10-08T09:15:14.552Z" }, + { url = "https://files.pythonhosted.org/packages/92/dc/c385f38f2c2433333345a82926c6bfa5ecfff3ef787201614317b58dd8be/msgpack-1.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:42eefe2c3e2af97ed470eec850facbe1b5ad1d6eacdbadc42ec98e7dcf68b4b7", size = 84315, upload-time = "2025-10-08T09:15:15.543Z" }, + { url = "https://files.pythonhosted.org/packages/d3/68/93180dce57f684a61a88a45ed13047558ded2be46f03acb8dec6d7c513af/msgpack-1.1.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1fdf7d83102bf09e7ce3357de96c59b627395352a4024f6e2458501f158bf999", size = 412721, upload-time = "2025-10-08T09:15:16.567Z" }, + { url = "https://files.pythonhosted.org/packages/5d/ba/459f18c16f2b3fc1a1ca871f72f07d70c07bf768ad0a507a698b8052ac58/msgpack-1.1.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fac4be746328f90caa3cd4bc67e6fe36ca2bf61d5c6eb6d895b6527e3f05071e", size = 424657, upload-time = "2025-10-08T09:15:17.825Z" }, + { url = "https://files.pythonhosted.org/packages/38/f8/4398c46863b093252fe67368b44edc6c13b17f4e6b0e4929dbf0bdb13f23/msgpack-1.1.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:fffee09044073e69f2bad787071aeec727183e7580443dfeb8556cbf1978d162", size = 402668, upload-time = "2025-10-08T09:15:19.003Z" }, + { url = "https://files.pythonhosted.org/packages/28/ce/698c1eff75626e4124b4d78e21cca0b4cc90043afb80a507626ea354ab52/msgpack-1.1.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5928604de9b032bc17f5099496417f113c45bc6bc21b5c6920caf34b3c428794", size = 419040, upload-time = "2025-10-08T09:15:20.183Z" }, + { url = "https://files.pythonhosted.org/packages/67/32/f3cd1667028424fa7001d82e10ee35386eea1408b93d399b09fb0aa7875f/msgpack-1.1.2-cp313-cp313-win32.whl", hash = "sha256:a7787d353595c7c7e145e2331abf8b7ff1e6673a6b974ded96e6d4ec09f00c8c", size = 65037, upload-time = "2025-10-08T09:15:21.416Z" }, + { url = "https://files.pythonhosted.org/packages/74/07/1ed8277f8653c40ebc65985180b007879f6a836c525b3885dcc6448ae6cb/msgpack-1.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:a465f0dceb8e13a487e54c07d04ae3ba131c7c5b95e2612596eafde1dccf64a9", size = 72631, upload-time = "2025-10-08T09:15:22.431Z" }, + { url = "https://files.pythonhosted.org/packages/e5/db/0314e4e2db56ebcf450f277904ffd84a7988b9e5da8d0d61ab2d057df2b6/msgpack-1.1.2-cp313-cp313-win_arm64.whl", hash = "sha256:e69b39f8c0aa5ec24b57737ebee40be647035158f14ed4b40e6f150077e21a84", size = 64118, upload-time = "2025-10-08T09:15:23.402Z" }, +] + [[package]] name = "mypy" version = "1.18.1" @@ -1102,6 +1137,7 @@ wheels = [ name = "power-grid-model-ds" source = { editable = "." } dependencies = [ + { name = "msgpack" }, { name = "numpy" }, { name = "power-grid-model" }, { name = "rustworkx" }, @@ -1143,7 +1179,8 @@ requires-dist = [ { name = "dash", marker = "extra == 'visualizer'", specifier = ">=3.0.0" }, { name = "dash-bootstrap-components", marker = "extra == 'visualizer'", specifier = ">=2.0.0" }, { name = "dash-cytoscape", marker = "extra == 'visualizer'", specifier = ">=1.0.2" }, - { name = "numpy", specifier = ">=1.21" }, + { name = "msgpack", specifier = ">=1.1.2" }, + { name = "numpy", specifier = ">=2.0" }, { name = "pandas", marker = "extra == 'pandas'", specifier = ">=2.2.1" }, { name = "power-grid-model", specifier = ">=1.7" }, { name = "rustworkx", specifier = ">=0.15.1" }, @@ -1572,6 +1609,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, ] +sdist = { url = "https://files.pythonhosted.org/packages/e7/b0/66d96f02120f79eeed86b5c5be04029b6821155f31ed4907a4e9f1460671/rustworkx-0.17.1.tar.gz", hash = "sha256:59ea01b4e603daffa4e8827316c1641eef18ae9032f0b1b14aa0181687e3108e", size = 399407, upload-time = "2025-09-15T16:29:46.429Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/20/24/8972ed631fa05fdec05a7bb7f1fc0f8e78ee761ab37e8a93d1ed396ba060/rustworkx-0.17.1-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c08fb8db041db052da404839b064ebfb47dcce04ba9a3e2eb79d0c65ab011da4", size = 2257491, upload-time = "2025-08-13T01:43:31.466Z" }, { url = "https://files.pythonhosted.org/packages/23/ae/7b6bbae5e0487ee42072dc6a46edf5db9731a0701ed648db22121fb7490c/rustworkx-0.17.1-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:4ef8e327dadf6500edd76fedb83f6d888b9266c58bcdbffd5a40c33835c9dd26", size = 2040175, upload-time = "2025-08-13T01:43:33.762Z" }, From d401ef6e6d7889a8bbadc9f88cbf5819de36109d Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Mon, 13 Oct 2025 19:55:57 +0200 Subject: [PATCH 04/19] chore: sonar feedback Signed-off-by: jaapschoutenalliander --- .../_core/utils/serialization.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py index 0c8de93..c377f42 100644 --- a/src/power_grid_model_ds/_core/utils/serialization.py +++ b/src/power_grid_model_ds/_core/utils/serialization.py @@ -9,7 +9,7 @@ import logging from ast import literal_eval from pathlib import Path -from typing import Dict, Optional +from typing import Dict, Literal, Optional import msgpack import numpy as np @@ -211,7 +211,7 @@ def _extract_msgpack_data(data: bytes, **kwargs): return input_data, extensions -def _get_serialization_path(path: Path, format_type: str = "auto") -> Path: +def _get_serialization_path(path: Path, format_type: Literal["json", "msgpack", "auto"] = "auto") -> Path: """Get the correct path for serialization format. Args: @@ -221,19 +221,22 @@ def _get_serialization_path(path: Path, format_type: str = "auto") -> Path: Returns: Path: Path with correct extension """ + JSON_EXTENSIONS = [".json"] + MSGPACK_EXTENSIONS = [".msgpack", ".mp"] + if format_type == "auto": - if path.suffix.lower() in [".json"]: + if path.suffix.lower() in JSON_EXTENSIONS: format_type = "json" - elif path.suffix.lower() in [".msgpack", ".mp"]: + elif path.suffix.lower() in MSGPACK_EXTENSIONS format_type = "msgpack" else: # Default to JSON format_type = "json" - if format_type == "json" and path.suffix.lower() != ".json": - return path.with_suffix(".json") - if format_type == "msgpack" and path.suffix.lower() not in [".msgpack", ".mp"]: - return path.with_suffix(".msgpack") + if format_type == "json" and path.suffix.lower() != JSON_EXTENSIONS[0]: + return path.with_suffix(JSON_EXTENSIONS[0]) + if format_type == "msgpack" and path.suffix.lower() not in MSGPACK_EXTENSIONS: + return path.with_suffix(MSGPACK_EXTENSIONS[0]) return path From f629e6a7bccf0b3ea91f7e952fc086bb986be189 Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Mon, 13 Oct 2025 20:06:14 +0200 Subject: [PATCH 05/19] chore: sonar feedback Signed-off-by: jaapschoutenalliander --- src/power_grid_model_ds/_core/utils/serialization.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py index c377f42..7f27b53 100644 --- a/src/power_grid_model_ds/_core/utils/serialization.py +++ b/src/power_grid_model_ds/_core/utils/serialization.py @@ -223,11 +223,11 @@ def _get_serialization_path(path: Path, format_type: Literal["json", "msgpack", """ JSON_EXTENSIONS = [".json"] MSGPACK_EXTENSIONS = [".msgpack", ".mp"] - + if format_type == "auto": if path.suffix.lower() in JSON_EXTENSIONS: format_type = "json" - elif path.suffix.lower() in MSGPACK_EXTENSIONS + elif path.suffix.lower() in MSGPACK_EXTENSIONS: format_type = "msgpack" else: # Default to JSON From b77f8e3969da9de88d9da8c2a7deeda82db2a903 Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Mon, 13 Oct 2025 20:09:22 +0200 Subject: [PATCH 06/19] chore: sonar feedback Signed-off-by: jaapschoutenalliander --- .../_core/utils/serialization.py | 16 ++++++++-------- tests/unit/utils/test_serialization.py | 5 ++++- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py index 7f27b53..54fcfff 100644 --- a/src/power_grid_model_ds/_core/utils/serialization.py +++ b/src/power_grid_model_ds/_core/utils/serialization.py @@ -221,22 +221,22 @@ def _get_serialization_path(path: Path, format_type: Literal["json", "msgpack", Returns: Path: Path with correct extension """ - JSON_EXTENSIONS = [".json"] - MSGPACK_EXTENSIONS = [".msgpack", ".mp"] + json_extensions = [".json"] + msgpack_extensions = [".msgpack", ".mp"] if format_type == "auto": - if path.suffix.lower() in JSON_EXTENSIONS: + if path.suffix.lower() in json_extensions: format_type = "json" - elif path.suffix.lower() in MSGPACK_EXTENSIONS: + elif path.suffix.lower() in msgpack_extensions: format_type = "msgpack" else: # Default to JSON format_type = "json" - if format_type == "json" and path.suffix.lower() != JSON_EXTENSIONS[0]: - return path.with_suffix(JSON_EXTENSIONS[0]) - if format_type == "msgpack" and path.suffix.lower() not in MSGPACK_EXTENSIONS: - return path.with_suffix(MSGPACK_EXTENSIONS[0]) + if format_type == "json" and path.suffix.lower() != json_extensions[0]: + return path.with_suffix(json_extensions[0]) + if format_type == "msgpack" and path.suffix.lower() not in msgpack_extensions: + return path.with_suffix(msgpack_extensions[0]) return path diff --git a/tests/unit/utils/test_serialization.py b/tests/unit/utils/test_serialization.py index 5e0ad61..14088bc 100644 --- a/tests/unit/utils/test_serialization.py +++ b/tests/unit/utils/test_serialization.py @@ -8,6 +8,7 @@ from dataclasses import dataclass from pathlib import Path from tempfile import TemporaryDirectory +from typing import Literal import numpy as np import pytest @@ -283,7 +284,9 @@ class TestUtilityFunctions: ("test.xyz", "msgpack", "test.msgpack"), ], ) - def test_serialization_path_handling(self, input_path: str, format_type: str, expected: str): + def test_serialization_path_handling( + self, input_path: str, format_type: Literal["json", "msgpack", "auto"], expected: str + ): """Test path handling and format detection""" result = _get_serialization_path(Path(input_path), format_type) assert result == Path(expected) From 108b9346981a58643faf98db6abd5ec6afdb3242 Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Wed, 22 Oct 2025 14:17:55 +0200 Subject: [PATCH 07/19] feat: remove msgpack Signed-off-by: jaapschoutenalliander --- pyproject.toml | 1 - src/power_grid_model_ds/__init__.py | 4 - .../_core/model/grids/base.py | 2 +- .../_core/utils/serialization.py | 121 +----------------- tests/unit/utils/test_serialization.py | 116 +++-------------- uv.lock | 37 ------ 6 files changed, 20 insertions(+), 261 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 110dab2..04b1aaf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,6 @@ dependencies = [ "power-grid-model>=1.7", "rustworkx>= 0.15.1", "numpy>=2.0", - "msgpack>=1.1.2", ] dynamic = ["version"] diff --git a/src/power_grid_model_ds/__init__.py b/src/power_grid_model_ds/__init__.py index 4f5c6aa..8638ffd 100644 --- a/src/power_grid_model_ds/__init__.py +++ b/src/power_grid_model_ds/__init__.py @@ -7,9 +7,7 @@ from power_grid_model_ds._core.model.grids.base import Grid from power_grid_model_ds._core.utils.serialization import ( load_grid_from_json, - load_grid_from_msgpack, save_grid_to_json, - save_grid_to_msgpack, ) __all__ = [ @@ -17,7 +15,5 @@ "GraphContainer", "PowerGridModelInterface", "save_grid_to_json", - "save_grid_to_msgpack", "load_grid_from_json", - "load_grid_from_msgpack", ] diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py index 35838f5..ed586d8 100644 --- a/src/power_grid_model_ds/_core/model/grids/base.py +++ b/src/power_grid_model_ds/_core/model/grids/base.py @@ -362,7 +362,7 @@ def get_downstream_nodes(self, node_id: int, inclusive: bool = False): def cache(self, cache_dir: Path, cache_name: str, compress: bool = True): """Cache Grid to a folder using pickle format. - Note: Consider using save_to_json() or save_to_msgpack() for better + Note: Consider using save_to_json() for better interoperability and standardized format. Args: diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py index 54fcfff..7edd43b 100644 --- a/src/power_grid_model_ds/_core/utils/serialization.py +++ b/src/power_grid_model_ds/_core/utils/serialization.py @@ -9,11 +9,10 @@ import logging from ast import literal_eval from pathlib import Path -from typing import Dict, Literal, Optional +from typing import Dict, Optional -import msgpack import numpy as np -from power_grid_model.utils import json_deserialize, json_serialize, msgpack_deserialize, msgpack_serialize +from power_grid_model.utils import json_deserialize, json_serialize from power_grid_model_ds._core.load_flow import PGM_ARRAYS, PowerGridModelInterface from power_grid_model_ds._core.model.arrays.base.array import FancyArray @@ -187,60 +186,6 @@ def _create_grid_from_input_data(input_data: Dict, target_grid_class=None): return interface.create_grid_from_input_data() -def _extract_msgpack_data(data: bytes, **kwargs): - """Extract input data and extensions from MessagePack data.""" - try: - data_dict = msgpack.unpackb(data, raw=False) - if isinstance(data_dict, dict) and EXTENSIONS_KEY in data_dict: - # Extract extensions and deserialize core data - extensions = data_dict.pop(EXTENSIONS_KEY, {}) - core_data = msgpack.packb(data_dict) - input_data = msgpack_deserialize(core_data, **kwargs) - else: - # No extensions, use power-grid-model directly - input_data = msgpack_deserialize(data, **kwargs) - extensions = {EXTENDED_COLUMNS_KEY: {}, CUSTOM_ARRAYS_KEY: {}} - except (msgpack.exceptions.ExtraData, ValueError, TypeError) as e: - # Handle MessagePack parsing failures: - # - ExtraData: malformed MessagePack data - # - ValueError/TypeError: invalid data structure or type issues - logger.warning(f"Failed to extract extensions from MessagePack data: {e}") - input_data = msgpack_deserialize(data, **kwargs) - extensions = {EXTENDED_COLUMNS_KEY: {}, CUSTOM_ARRAYS_KEY: {}} - - return input_data, extensions - - -def _get_serialization_path(path: Path, format_type: Literal["json", "msgpack", "auto"] = "auto") -> Path: - """Get the correct path for serialization format. - - Args: - path: Base path - format_type: "json", "msgpack", or "auto" to detect from extension - - Returns: - Path: Path with correct extension - """ - json_extensions = [".json"] - msgpack_extensions = [".msgpack", ".mp"] - - if format_type == "auto": - if path.suffix.lower() in json_extensions: - format_type = "json" - elif path.suffix.lower() in msgpack_extensions: - format_type = "msgpack" - else: - # Default to JSON - format_type = "json" - - if format_type == "json" and path.suffix.lower() != json_extensions[0]: - return path.with_suffix(json_extensions[0]) - if format_type == "msgpack" and path.suffix.lower() not in msgpack_extensions: - return path.with_suffix(msgpack_extensions[0]) - - return path - - def save_grid_to_json( grid, path: Path, @@ -303,65 +248,3 @@ def load_grid_from_json(path: Path, target_grid_class=None): _restore_extensions_data(grid, extensions) return grid - - -def save_grid_to_msgpack(grid, path: Path, use_compact_list: bool = True, preserve_extensions: bool = True) -> Path: - """Save a Grid object to MessagePack format with extensions support. - - Args: - grid: The Grid object to serialize - path: The file path to save to - use_compact_list: Whether to use compact list format - preserve_extensions: Whether to save extended columns and custom arrays - - Returns: - Path: The path where the file was saved - """ - path.parent.mkdir(parents=True, exist_ok=True) - - # Convert Grid to power-grid-model input format and serialize - interface = PowerGridModelInterface(grid=grid) - input_data = interface.create_input_from_grid() - - core_data = msgpack_serialize(input_data, use_compact_list=use_compact_list) - - # Add extensions if requested (requires re-serialization for MessagePack) - if preserve_extensions: - extensions = _extract_extensions_data(grid) - if extensions[EXTENDED_COLUMNS_KEY] or extensions[CUSTOM_ARRAYS_KEY]: - core_dict = msgpack.unpackb(core_data, raw=False) - core_dict[EXTENSIONS_KEY] = extensions - serialized_data = msgpack.packb(core_dict) - else: - serialized_data = core_data - else: - serialized_data = core_data - - # Write to file - with open(path, "wb") as f: - f.write(serialized_data) - - return path - - -def load_grid_from_msgpack(path: Path, target_grid_class=None): - """Load a Grid object from MessagePack format with cross-type loading support. - - Args: - path: The file path to load from - target_grid_class: Optional Grid class to load into. If None, uses default Grid. - - Returns: - Grid: The deserialized Grid object of the specified target class - """ - with open(path, "rb") as f: - data = f.read() - - # Extract extensions and deserialize core data - input_data, extensions = _extract_msgpack_data(data) - - # Create grid and restore extensions - grid = _create_grid_from_input_data(input_data, target_grid_class) - _restore_extensions_data(grid, extensions) - - return grid diff --git a/tests/unit/utils/test_serialization.py b/tests/unit/utils/test_serialization.py index 14088bc..ecc074e 100644 --- a/tests/unit/utils/test_serialization.py +++ b/tests/unit/utils/test_serialization.py @@ -8,7 +8,6 @@ from dataclasses import dataclass from pathlib import Path from tempfile import TemporaryDirectory -from typing import Literal import numpy as np import pytest @@ -18,12 +17,9 @@ from power_grid_model_ds._core.model.arrays.base.array import FancyArray from power_grid_model_ds._core.utils.serialization import ( _extract_extensions_data, - _get_serialization_path, _restore_extensions_data, load_grid_from_json, - load_grid_from_msgpack, save_grid_to_json, - save_grid_to_msgpack, ) from power_grid_model_ds.arrays import LineArray from power_grid_model_ds.arrays import NodeArray as BaseNodeArray @@ -95,46 +91,25 @@ def extended_grid(): class TestSerializationFormats: """Test serialization across different formats and configurations""" - @pytest.mark.parametrize( - "format_type,preserve_ext", [("json", True), ("json", False), ("msgpack", True), ("msgpack", False)] - ) - def test_basic_serialization_roundtrip( - self, basic_grid: Grid, temp_dir: Path, format_type: str, preserve_ext: bool - ): + @pytest.mark.parametrize("preserve_ext", [(True), (False)]) + def test_basic_serialization_roundtrip(self, basic_grid: Grid, temp_dir: Path, preserve_ext: bool): """Test basic serialization roundtrip for all formats""" - ext = "json" if format_type == "json" else "msgpack" - path = temp_dir / f"test.{ext}" - - # Save - if format_type == "json": - result_path = save_grid_to_json(basic_grid, path, preserve_extensions=preserve_ext) - else: - result_path = save_grid_to_msgpack(basic_grid, path, preserve_extensions=preserve_ext) - + path = temp_dir / "test.json" + result_path = save_grid_to_json(basic_grid, path, preserve_extensions=preserve_ext) assert result_path.exists() # Load and verify - if format_type == "json": - loaded_grid = load_grid_from_json(path, target_grid_class=Grid) - else: - loaded_grid = load_grid_from_msgpack(path, target_grid_class=Grid) + loaded_grid = load_grid_from_json(path, target_grid_class=Grid) assert loaded_grid.node.size == basic_grid.node.size assert loaded_grid.line.size == basic_grid.line.size assert list(loaded_grid.node.id) == list(basic_grid.node.id) - @pytest.mark.parametrize("format_type", ["json", "msgpack"]) - def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, temp_dir: Path, format_type: str): + def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, temp_dir: Path): """Test extended serialization preserving custom data""" - ext = "json" if format_type == "json" else "msgpack" - path = temp_dir / f"extended.{ext}" + path = temp_dir / "extended.json" - # Save with extensions - if format_type == "json": - save_grid_to_json(extended_grid, path, preserve_extensions=True) - loaded_grid = load_grid_from_json(path, target_grid_class=ExtendedGrid) - else: - save_grid_to_msgpack(extended_grid, path, preserve_extensions=True) - loaded_grid = load_grid_from_msgpack(path, target_grid_class=ExtendedGrid) + save_grid_to_json(extended_grid, path, preserve_extensions=True) + loaded_grid = load_grid_from_json(path, target_grid_class=ExtendedGrid) # Verify core data assert loaded_grid.node.size == extended_grid.node.size @@ -148,37 +123,25 @@ def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, tem class TestCrossTypeCompatibility: """Test cross-type loading and compatibility""" - @pytest.mark.parametrize("format_type", ["json", "msgpack"]) - def test_basic_to_extended_loading(self, basic_grid: Grid, temp_dir: Path, format_type: str): + def test_basic_to_extended_loading(self, basic_grid: Grid, temp_dir: Path): """Test loading basic grid into extended type""" - ext = "json" if format_type == "json" else "msgpack" - path = temp_dir / f"basic.{ext}" + path = temp_dir / "basic.json" # Save basic grid - if format_type == "json": - save_grid_to_json(basic_grid, path) - loaded_grid = load_grid_from_json(path, target_grid_class=ExtendedGrid) - else: - save_grid_to_msgpack(basic_grid, path) - loaded_grid = load_grid_from_msgpack(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 assert loaded_grid.node.size == basic_grid.node.size assert loaded_grid.line.size == basic_grid.line.size - @pytest.mark.parametrize("format_type", ["json", "msgpack"]) - def test_extended_to_basic_loading(self, extended_grid: ExtendedGrid, temp_dir: Path, format_type: str): + def test_extended_to_basic_loading(self, extended_grid: ExtendedGrid, temp_dir: Path): """Test loading extended grid into basic type""" - ext = "json" if format_type == "json" else "msgpack" - path = temp_dir / f"extended.{ext}" + path = temp_dir / "extended.json" # Save extended grid - if format_type == "json": - save_grid_to_json(extended_grid, path, preserve_extensions=True) - loaded_grid = load_grid_from_json(path, target_grid_class=Grid) - else: - save_grid_to_msgpack(extended_grid, path, preserve_extensions=True) - loaded_grid = load_grid_from_msgpack(path, target_grid_class=Grid) + save_grid_to_json(extended_grid, path, preserve_extensions=True) + loaded_grid = load_grid_from_json(path, target_grid_class=Grid) # Core data should transfer assert loaded_grid.node.size == extended_grid.node.size @@ -251,46 +214,6 @@ class GridWithCustomArray(Grid): np.testing.assert_array_almost_equal(loaded_grid.custom_metadata.metadata_value, [1.5, 2.5, 3.5]) np.testing.assert_array_equal(loaded_grid.custom_metadata.category, [1, 2, 1]) - # Test MessagePack serialization - msgpack_path = temp_dir / "custom_array.msgpack" - save_grid_to_msgpack(grid, msgpack_path, preserve_extensions=True) - - # Load back and verify - loaded_grid_mp = load_grid_from_msgpack(msgpack_path, target_grid_class=GridWithCustomArray) - - # Verify core data - assert loaded_grid_mp.node.size == 2 - np.testing.assert_array_equal(loaded_grid_mp.node.id, [1, 2]) - - # Verify custom array was preserved - assert hasattr(loaded_grid_mp, "custom_metadata") - assert loaded_grid_mp.custom_metadata.size == 3 - np.testing.assert_array_equal(loaded_grid_mp.custom_metadata.id, [100, 200, 300]) - np.testing.assert_array_almost_equal(loaded_grid_mp.custom_metadata.metadata_value, [1.5, 2.5, 3.5]) - np.testing.assert_array_equal(loaded_grid_mp.custom_metadata.category, [1, 2, 1]) - - -class TestUtilityFunctions: - """Test utility functions and path handling""" - - @pytest.mark.parametrize( - "input_path,format_type,expected", - [ - ("test.json", "auto", "test.json"), - ("test.msgpack", "auto", "test.msgpack"), - ("test.mp", "auto", "test.mp"), - ("test.xyz", "auto", "test.json"), # Unknown defaults to JSON - ("test.xyz", "json", "test.json"), - ("test.xyz", "msgpack", "test.msgpack"), - ], - ) - def test_serialization_path_handling( - self, input_path: str, format_type: Literal["json", "msgpack", "auto"], expected: str - ): - """Test path handling and format detection""" - result = _get_serialization_path(Path(input_path), format_type) - assert result == Path(expected) - class TestSpecialCases: """Test special cases and edge scenarios""" @@ -300,18 +223,13 @@ def test_empty_grid_handling(self, temp_dir: Path): empty_grid = Grid.empty() json_path = temp_dir / "empty.json" - msgpack_path = temp_dir / "empty.msgpack" # Should handle empty grids save_grid_to_json(empty_grid, json_path) - save_grid_to_msgpack(empty_grid, msgpack_path) # Should load back as empty loaded_json = load_grid_from_json(json_path, target_grid_class=Grid) - loaded_msgpack = load_grid_from_msgpack(msgpack_path, target_grid_class=Grid) - assert loaded_json.node.size == 0 - assert loaded_msgpack.node.size == 0 def test_custom_array_extraction_edge_cases(self, temp_dir: Path): """Test edge cases in custom array extraction""" diff --git a/uv.lock b/uv.lock index 5346f13..c36f61d 100644 --- a/uv.lock +++ b/uv.lock @@ -763,41 +763,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] -[[package]] -name = "msgpack" -version = "1.1.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/4d/f2/bfb55a6236ed8725a96b0aa3acbd0ec17588e6a2c3b62a93eb513ed8783f/msgpack-1.1.2.tar.gz", hash = "sha256:3b60763c1373dd60f398488069bcdc703cd08a711477b5d480eecc9f9626f47e", size = 173581, upload-time = "2025-10-08T09:15:56.596Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2c/97/560d11202bcd537abca693fd85d81cebe2107ba17301de42b01ac1677b69/msgpack-1.1.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:2e86a607e558d22985d856948c12a3fa7b42efad264dca8a3ebbcfa2735d786c", size = 82271, upload-time = "2025-10-08T09:14:49.967Z" }, - { url = "https://files.pythonhosted.org/packages/83/04/28a41024ccbd67467380b6fb440ae916c1e4f25e2cd4c63abe6835ac566e/msgpack-1.1.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:283ae72fc89da59aa004ba147e8fc2f766647b1251500182fac0350d8af299c0", size = 84914, upload-time = "2025-10-08T09:14:50.958Z" }, - { url = "https://files.pythonhosted.org/packages/71/46/b817349db6886d79e57a966346cf0902a426375aadc1e8e7a86a75e22f19/msgpack-1.1.2-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:61c8aa3bd513d87c72ed0b37b53dd5c5a0f58f2ff9f26e1555d3bd7948fb7296", size = 416962, upload-time = "2025-10-08T09:14:51.997Z" }, - { url = "https://files.pythonhosted.org/packages/da/e0/6cc2e852837cd6086fe7d8406af4294e66827a60a4cf60b86575a4a65ca8/msgpack-1.1.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:454e29e186285d2ebe65be34629fa0e8605202c60fbc7c4c650ccd41870896ef", size = 426183, upload-time = "2025-10-08T09:14:53.477Z" }, - { url = "https://files.pythonhosted.org/packages/25/98/6a19f030b3d2ea906696cedd1eb251708e50a5891d0978b012cb6107234c/msgpack-1.1.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:7bc8813f88417599564fafa59fd6f95be417179f76b40325b500b3c98409757c", size = 411454, upload-time = "2025-10-08T09:14:54.648Z" }, - { url = "https://files.pythonhosted.org/packages/b7/cd/9098fcb6adb32187a70b7ecaabf6339da50553351558f37600e53a4a2a23/msgpack-1.1.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:bafca952dc13907bdfdedfc6a5f579bf4f292bdd506fadb38389afa3ac5b208e", size = 422341, upload-time = "2025-10-08T09:14:56.328Z" }, - { url = "https://files.pythonhosted.org/packages/e6/ae/270cecbcf36c1dc85ec086b33a51a4d7d08fc4f404bdbc15b582255d05ff/msgpack-1.1.2-cp311-cp311-win32.whl", hash = "sha256:602b6740e95ffc55bfb078172d279de3773d7b7db1f703b2f1323566b878b90e", size = 64747, upload-time = "2025-10-08T09:14:57.882Z" }, - { url = "https://files.pythonhosted.org/packages/2a/79/309d0e637f6f37e83c711f547308b91af02b72d2326ddd860b966080ef29/msgpack-1.1.2-cp311-cp311-win_amd64.whl", hash = "sha256:d198d275222dc54244bf3327eb8cbe00307d220241d9cec4d306d49a44e85f68", size = 71633, upload-time = "2025-10-08T09:14:59.177Z" }, - { url = "https://files.pythonhosted.org/packages/73/4d/7c4e2b3d9b1106cd0aa6cb56cc57c6267f59fa8bfab7d91df5adc802c847/msgpack-1.1.2-cp311-cp311-win_arm64.whl", hash = "sha256:86f8136dfa5c116365a8a651a7d7484b65b13339731dd6faebb9a0242151c406", size = 64755, upload-time = "2025-10-08T09:15:00.48Z" }, - { url = "https://files.pythonhosted.org/packages/ad/bd/8b0d01c756203fbab65d265859749860682ccd2a59594609aeec3a144efa/msgpack-1.1.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:70a0dff9d1f8da25179ffcf880e10cf1aad55fdb63cd59c9a49a1b82290062aa", size = 81939, upload-time = "2025-10-08T09:15:01.472Z" }, - { url = "https://files.pythonhosted.org/packages/34/68/ba4f155f793a74c1483d4bdef136e1023f7bcba557f0db4ef3db3c665cf1/msgpack-1.1.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:446abdd8b94b55c800ac34b102dffd2f6aa0ce643c55dfc017ad89347db3dbdb", size = 85064, upload-time = "2025-10-08T09:15:03.764Z" }, - { url = "https://files.pythonhosted.org/packages/f2/60/a064b0345fc36c4c3d2c743c82d9100c40388d77f0b48b2f04d6041dbec1/msgpack-1.1.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c63eea553c69ab05b6747901b97d620bb2a690633c77f23feb0c6a947a8a7b8f", size = 417131, upload-time = "2025-10-08T09:15:05.136Z" }, - { url = "https://files.pythonhosted.org/packages/65/92/a5100f7185a800a5d29f8d14041f61475b9de465ffcc0f3b9fba606e4505/msgpack-1.1.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:372839311ccf6bdaf39b00b61288e0557916c3729529b301c52c2d88842add42", size = 427556, upload-time = "2025-10-08T09:15:06.837Z" }, - { url = "https://files.pythonhosted.org/packages/f5/87/ffe21d1bf7d9991354ad93949286f643b2bb6ddbeab66373922b44c3b8cc/msgpack-1.1.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2929af52106ca73fcb28576218476ffbb531a036c2adbcf54a3664de124303e9", size = 404920, upload-time = "2025-10-08T09:15:08.179Z" }, - { url = "https://files.pythonhosted.org/packages/ff/41/8543ed2b8604f7c0d89ce066f42007faac1eaa7d79a81555f206a5cdb889/msgpack-1.1.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:be52a8fc79e45b0364210eef5234a7cf8d330836d0a64dfbb878efa903d84620", size = 415013, upload-time = "2025-10-08T09:15:09.83Z" }, - { url = "https://files.pythonhosted.org/packages/41/0d/2ddfaa8b7e1cee6c490d46cb0a39742b19e2481600a7a0e96537e9c22f43/msgpack-1.1.2-cp312-cp312-win32.whl", hash = "sha256:1fff3d825d7859ac888b0fbda39a42d59193543920eda9d9bea44d958a878029", size = 65096, upload-time = "2025-10-08T09:15:11.11Z" }, - { url = "https://files.pythonhosted.org/packages/8c/ec/d431eb7941fb55a31dd6ca3404d41fbb52d99172df2e7707754488390910/msgpack-1.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:1de460f0403172cff81169a30b9a92b260cb809c4cb7e2fc79ae8d0510c78b6b", size = 72708, upload-time = "2025-10-08T09:15:12.554Z" }, - { url = "https://files.pythonhosted.org/packages/c5/31/5b1a1f70eb0e87d1678e9624908f86317787b536060641d6798e3cf70ace/msgpack-1.1.2-cp312-cp312-win_arm64.whl", hash = "sha256:be5980f3ee0e6bd44f3a9e9dea01054f175b50c3e6cdb692bc9424c0bbb8bf69", size = 64119, upload-time = "2025-10-08T09:15:13.589Z" }, - { url = "https://files.pythonhosted.org/packages/6b/31/b46518ecc604d7edf3a4f94cb3bf021fc62aa301f0cb849936968164ef23/msgpack-1.1.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:4efd7b5979ccb539c221a4c4e16aac1a533efc97f3b759bb5a5ac9f6d10383bf", size = 81212, upload-time = "2025-10-08T09:15:14.552Z" }, - { url = "https://files.pythonhosted.org/packages/92/dc/c385f38f2c2433333345a82926c6bfa5ecfff3ef787201614317b58dd8be/msgpack-1.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:42eefe2c3e2af97ed470eec850facbe1b5ad1d6eacdbadc42ec98e7dcf68b4b7", size = 84315, upload-time = "2025-10-08T09:15:15.543Z" }, - { url = "https://files.pythonhosted.org/packages/d3/68/93180dce57f684a61a88a45ed13047558ded2be46f03acb8dec6d7c513af/msgpack-1.1.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1fdf7d83102bf09e7ce3357de96c59b627395352a4024f6e2458501f158bf999", size = 412721, upload-time = "2025-10-08T09:15:16.567Z" }, - { url = "https://files.pythonhosted.org/packages/5d/ba/459f18c16f2b3fc1a1ca871f72f07d70c07bf768ad0a507a698b8052ac58/msgpack-1.1.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fac4be746328f90caa3cd4bc67e6fe36ca2bf61d5c6eb6d895b6527e3f05071e", size = 424657, upload-time = "2025-10-08T09:15:17.825Z" }, - { url = "https://files.pythonhosted.org/packages/38/f8/4398c46863b093252fe67368b44edc6c13b17f4e6b0e4929dbf0bdb13f23/msgpack-1.1.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:fffee09044073e69f2bad787071aeec727183e7580443dfeb8556cbf1978d162", size = 402668, upload-time = "2025-10-08T09:15:19.003Z" }, - { url = "https://files.pythonhosted.org/packages/28/ce/698c1eff75626e4124b4d78e21cca0b4cc90043afb80a507626ea354ab52/msgpack-1.1.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5928604de9b032bc17f5099496417f113c45bc6bc21b5c6920caf34b3c428794", size = 419040, upload-time = "2025-10-08T09:15:20.183Z" }, - { url = "https://files.pythonhosted.org/packages/67/32/f3cd1667028424fa7001d82e10ee35386eea1408b93d399b09fb0aa7875f/msgpack-1.1.2-cp313-cp313-win32.whl", hash = "sha256:a7787d353595c7c7e145e2331abf8b7ff1e6673a6b974ded96e6d4ec09f00c8c", size = 65037, upload-time = "2025-10-08T09:15:21.416Z" }, - { url = "https://files.pythonhosted.org/packages/74/07/1ed8277f8653c40ebc65985180b007879f6a836c525b3885dcc6448ae6cb/msgpack-1.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:a465f0dceb8e13a487e54c07d04ae3ba131c7c5b95e2612596eafde1dccf64a9", size = 72631, upload-time = "2025-10-08T09:15:22.431Z" }, - { url = "https://files.pythonhosted.org/packages/e5/db/0314e4e2db56ebcf450f277904ffd84a7988b9e5da8d0d61ab2d057df2b6/msgpack-1.1.2-cp313-cp313-win_arm64.whl", hash = "sha256:e69b39f8c0aa5ec24b57737ebee40be647035158f14ed4b40e6f150077e21a84", size = 64118, upload-time = "2025-10-08T09:15:23.402Z" }, -] - [[package]] name = "mypy" version = "1.18.1" @@ -1137,7 +1102,6 @@ wheels = [ name = "power-grid-model-ds" source = { editable = "." } dependencies = [ - { name = "msgpack" }, { name = "numpy" }, { name = "power-grid-model" }, { name = "rustworkx" }, @@ -1179,7 +1143,6 @@ requires-dist = [ { name = "dash", marker = "extra == 'visualizer'", specifier = ">=3.0.0" }, { name = "dash-bootstrap-components", marker = "extra == 'visualizer'", specifier = ">=2.0.0" }, { name = "dash-cytoscape", marker = "extra == 'visualizer'", specifier = ">=1.0.2" }, - { name = "msgpack", specifier = ">=1.1.2" }, { name = "numpy", specifier = ">=2.0" }, { name = "pandas", marker = "extra == 'pandas'", specifier = ">=2.2.1" }, { name = "power-grid-model", specifier = ">=1.7" }, From a6f0cbe74419278981f94ae5b0940fe3dbbe86f3 Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Wed, 22 Oct 2025 15:11:07 +0200 Subject: [PATCH 08/19] feat: bypass pgm json conversion for simplicity Signed-off-by: jaapschoutenalliander --- .../_core/utils/serialization.py | 201 +++--------------- tests/unit/utils/test_serialization.py | 76 +++---- 2 files changed, 54 insertions(+), 223 deletions(-) diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py index 7edd43b..c91b307 100644 --- a/src/power_grid_model_ds/_core/utils/serialization.py +++ b/src/power_grid_model_ds/_core/utils/serialization.py @@ -7,155 +7,27 @@ import dataclasses import json import logging -from ast import literal_eval from pathlib import Path from typing import Dict, Optional -import numpy as np -from power_grid_model.utils import json_deserialize, json_serialize - -from power_grid_model_ds._core.load_flow import PGM_ARRAYS, PowerGridModelInterface from power_grid_model_ds._core.model.arrays.base.array import FancyArray - -# Constants -EXTENDED_COLUMNS_KEY = "extended_columns" -CUSTOM_ARRAYS_KEY = "custom_arrays" -EXTENSIONS_KEY = "pgm_ds_extensions" +from power_grid_model_ds._core.model.grids.base import Grid logger = logging.getLogger(__name__) -def _extract_extensions_data(grid) -> Dict[str, Dict]: - """Extract extended columns and non-PGM arrays from a Grid object. - - Args: - grid: The Grid object - - Returns: - Dict containing extensions data with keys EXTENDED_COLUMNS_KEY and CUSTOM_ARRAYS_KEY - """ - extensions: dict = {EXTENDED_COLUMNS_KEY: {}, CUSTOM_ARRAYS_KEY: {}} - - for field in dataclasses.fields(grid): - if field.name in ["graphs", "_id_counter"]: - continue - - array = getattr(grid, field.name) - if not isinstance(array, FancyArray) or array.size == 0: - continue - - array_name = field.name - - if array_name in PGM_ARRAYS: - # Extract extended columns for PGM arrays - _extract_extended_columns(grid, array_name, array, extensions) - else: - # Store custom arrays not in PGM_ARRAYS - extensions[CUSTOM_ARRAYS_KEY][array_name] = {"dtype": str(array.dtype), "data": array.data.tolist()} - - return extensions - - -def _extract_extended_columns(grid, array_name: str, array: FancyArray, extensions: Dict) -> None: - """Extract extended columns from a PGM array.""" - try: - interface = PowerGridModelInterface(grid=grid) - # pylint: disable=protected-access # Accessing internal method for extension extraction - pgm_array = interface._create_power_grid_array(array_name) - pgm_columns = set(pgm_array.dtype.names or []) - ds_columns = set(array.columns) - - # Find extended columns (columns in DS but not in PGM) - extended_cols = ds_columns - pgm_columns - if extended_cols: - extensions[EXTENDED_COLUMNS_KEY][array_name] = {col: array[col].tolist() for col in extended_cols} - except (AttributeError, KeyError, TypeError, ValueError) as e: - # Handle various failure modes: - # - KeyError: array_name not found in PGM arrays - # - AttributeError: array missing dtype/columns or interface method missing - # - TypeError/ValueError: invalid array configuration or data conversion issues - logger.warning(f"Failed to extract extensions for array '{array_name}': {e}") - extensions[CUSTOM_ARRAYS_KEY][array_name] = {"dtype": str(array.dtype), "data": array.data.tolist()} - - -def _restore_extensions_data(grid, extensions_data: Dict) -> None: - """Restore extended columns and custom arrays to a Grid object. - - Args: - grid: The Grid object to restore extensions to - extensions_data: Extensions data from _extract_extensions_data - """ - # Restore extended columns - _restore_extended_columns(grid, extensions_data.get(EXTENDED_COLUMNS_KEY, {})) - - # Restore custom arrays - _restore_custom_arrays(grid, extensions_data.get(CUSTOM_ARRAYS_KEY, {})) - - -def _restore_extended_columns(grid, extended_columns: Dict) -> None: - """Restore extended columns to existing arrays.""" - for array_name, extended_cols in extended_columns.items(): - if not hasattr(grid, array_name): - logger.warning(f"Grid has no attribute '{array_name}' to restore") - continue - - array = getattr(grid, array_name) - if not isinstance(array, FancyArray) or array.size == 0: - continue - - for col_name, values in extended_cols.items(): - # if hasattr(array, col_name): - try: - array[col_name] = values - except (AttributeError, IndexError, ValueError, TypeError) as e: - # Handle assignment failures: - # - IndexError: array size mismatch - # - ValueError/TypeError: incompatible data types - # - AttributeError: array doesn't support assignment - logger.warning(f"Failed to restore column '{col_name}' in array '{array_name}': {e}") - - -def _parse_dtype(dtype_str: str) -> np.dtype: - """Parse a dtype string into a numpy dtype.""" - if not isinstance(dtype_str, str): - raise ValueError(f"Invalid dtype string: {dtype_str}") - - # Use numpy's dtype parsing - handle both eval-style and direct strings - if dtype_str.startswith("dtype("): - clean_dtype_str = dtype_str.replace("dtype(", "").replace(")", "") - else: - clean_dtype_str = dtype_str - - # Use eval for complex dtype strings like "[('field', 'type'), ...]" - if clean_dtype_str.startswith("[") and clean_dtype_str.endswith("]"): - return np.dtype(literal_eval(clean_dtype_str)) - return np.dtype(clean_dtype_str) - - -def _construct_numpy_from_list(raw_data, dtype: np.dtype) -> np.ndarray: - """Construct a numpy array from a list with the specified dtype.""" - if dtype.names: # Structured dtype - # Convert from list of lists to list of tuples for structured array - if isinstance(raw_data[0], (list, tuple)) and len(raw_data[0]) == len(dtype.names): - data = np.array([tuple(row) for row in raw_data], dtype=dtype) - else: - data = np.array(raw_data, dtype=dtype) - else: - data = np.array(raw_data, dtype=dtype) - return data - - -def _restore_custom_arrays(grid, custom_arrays: Dict) -> None: +def _restore_grid_arrays(grid, custom_arrays: Dict) -> None: """Restore custom arrays to the grid.""" for array_name, array_info in custom_arrays.items(): if not hasattr(grid, array_name): continue try: - dtype = _parse_dtype(dtype_str=array_info["dtype"]) - data = _construct_numpy_from_list(array_info["data"], dtype) array_field = grid.find_array_field(getattr(grid, array_name).__class__) - restored_array = array_field.type(data=data) + matched_columns = { + col: array_info["data"][col] for col in array_field.type().columns if col in array_info["data"] + } + restored_array = array_field.type(**matched_columns) setattr(grid, array_name, restored_array) except (AttributeError, KeyError, ValueError, TypeError) as e: # Handle restoration failures: @@ -165,59 +37,35 @@ def _restore_custom_arrays(grid, custom_arrays: Dict) -> None: logger.warning(f"Failed to restore custom array '{array_name}': {e}") -def _create_grid_from_input_data(input_data: Dict, target_grid_class=None): - """Create a Grid object from power-grid-model input data. - - Args: - input_data: Power-grid-model input data - target_grid_class: Optional Grid class to create. If None, uses default Grid. - - Returns: - Grid object populated with the input data - """ - if target_grid_class is not None: - # Create empty grid of target type and populate it with input data - target_grid = target_grid_class.empty() - interface = PowerGridModelInterface(grid=target_grid, input_data=input_data) - return interface.create_grid_from_input_data() - - # Use default Grid type - interface = PowerGridModelInterface(input_data=input_data) - return interface.create_grid_from_input_data() - - def save_grid_to_json( grid, path: Path, - use_compact_list: bool = True, indent: Optional[int] = None, - preserve_extensions: bool = True, ) -> 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 - use_compact_list: Whether to use compact list format indent: JSON indentation (None for compact, positive int for indentation) - preserve_extensions: Whether to save extended columns and custom arrays Returns: Path: The path where the file was saved """ path.parent.mkdir(parents=True, exist_ok=True) - # Convert Grid to power-grid-model input format and serialize - interface = PowerGridModelInterface(grid=grid) - input_data = interface.create_input_from_grid() + serialized_data = {} + for field in dataclasses.fields(grid): + if field.name in ["graphs", "_id_counter"]: + continue - core_data = json_serialize(input_data, use_compact_list=use_compact_list) + array = getattr(grid, field.name) + if not isinstance(array, FancyArray) or array.size == 0: + continue - # Parse and add extensions if requested - serialized_data = json.loads(core_data) - if preserve_extensions: - extensions = _extract_extensions_data(grid) - if extensions[EXTENDED_COLUMNS_KEY] or extensions[CUSTOM_ARRAYS_KEY]: - serialized_data[EXTENSIONS_KEY] = extensions + array_name = field.name + serialized_data[array_name] = { + "data": {name: array[name].tolist() for name in array.dtype.names}, + } # Write to file with open(path, "w", encoding="utf-8") as f: @@ -237,14 +85,13 @@ def load_grid_from_json(path: Path, target_grid_class=None): Grid: The deserialized Grid object of the specified target class """ with open(path, "r", encoding="utf-8") as f: - data = json.load(f) + input_data = json.load(f) - # Extract extensions and deserialize core data - extensions = data.pop(EXTENSIONS_KEY, {EXTENDED_COLUMNS_KEY: {}, CUSTOM_ARRAYS_KEY: {}}) - input_data = json_deserialize(json.dumps(data)) + if target_grid_class is None: + target_grid = Grid.empty() + else: + target_grid = target_grid_class.empty() - # Create grid and restore extensions - grid = _create_grid_from_input_data(input_data, target_grid_class) - _restore_extensions_data(grid, extensions) + _restore_grid_arrays(target_grid, input_data) - return grid + return target_grid diff --git a/tests/unit/utils/test_serialization.py b/tests/unit/utils/test_serialization.py index ecc074e..9f353d8 100644 --- a/tests/unit/utils/test_serialization.py +++ b/tests/unit/utils/test_serialization.py @@ -16,8 +16,6 @@ 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 ( - _extract_extensions_data, - _restore_extensions_data, load_grid_from_json, save_grid_to_json, ) @@ -48,6 +46,9 @@ class ExtendedGrid(Grid): node: ExtendedNodeArray line: ExtendedLineArray + # value_extension: float = 0.0 + # dict_extension: dict = dict() + @pytest.fixture def temp_dir(): @@ -91,11 +92,10 @@ def extended_grid(): class TestSerializationFormats: """Test serialization across different formats and configurations""" - @pytest.mark.parametrize("preserve_ext", [(True), (False)]) - def test_basic_serialization_roundtrip(self, basic_grid: Grid, temp_dir: Path, preserve_ext: bool): + def test_basic_serialization_roundtrip(self, basic_grid: Grid, temp_dir: Path): """Test basic serialization roundtrip for all formats""" path = temp_dir / "test.json" - result_path = save_grid_to_json(basic_grid, path, preserve_extensions=preserve_ext) + result_path = save_grid_to_json(basic_grid, path) assert result_path.exists() # Load and verify @@ -108,7 +108,7 @@ def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, tem """Test extended serialization preserving custom data""" path = temp_dir / "extended.json" - save_grid_to_json(extended_grid, path, preserve_extensions=True) + save_grid_to_json(extended_grid, path) loaded_grid = load_grid_from_json(path, target_grid_class=ExtendedGrid) # Verify core data @@ -140,7 +140,7 @@ def test_extended_to_basic_loading(self, extended_grid: ExtendedGrid, temp_dir: path = temp_dir / "extended.json" # Save extended grid - save_grid_to_json(extended_grid, path, preserve_extensions=True) + save_grid_to_json(extended_grid, path) loaded_grid = load_grid_from_json(path, target_grid_class=Grid) # Core data should transfer @@ -151,22 +151,6 @@ def test_extended_to_basic_loading(self, extended_grid: ExtendedGrid, temp_dir: class TestExtensionHandling: """Test extension data handling and edge cases""" - def test_missing_extension_keys(self): - """Test graceful handling of missing extension keys""" - basic_grid = Grid.empty() - - # Test various malformed extension data - test_cases = [ - {}, # Empty - {"extended_columns": {}}, # Missing custom_arrays - {"custom_arrays": {}}, # Missing extended_columns - {"extended_columns": {"test": "value"}}, # Invalid structure - ] - - for extensions in test_cases: - # Should not raise - _restore_extensions_data(basic_grid, extensions) - def test_custom_array_serialization_roundtrip(self, temp_dir: Path): """Test serialization and loading of grids with custom arrays""" @@ -198,7 +182,7 @@ class GridWithCustomArray(Grid): # Test JSON serialization json_path = temp_dir / "custom_array.json" - save_grid_to_json(grid, json_path, preserve_extensions=True) + save_grid_to_json(grid, json_path) # Load back and verify loaded_grid = load_grid_from_json(json_path, target_grid_class=GridWithCustomArray) @@ -231,31 +215,31 @@ def test_empty_grid_handling(self, temp_dir: Path): loaded_json = load_grid_from_json(json_path, target_grid_class=Grid) assert loaded_json.node.size == 0 - def test_custom_array_extraction_edge_cases(self, temp_dir: Path): - """Test edge cases in custom array extraction""" - # Test with grid that has complex custom arrays that might cause extraction issues - extended_grid = ExtendedGrid.empty() + # def test_custom_array_extraction_edge_cases(self, temp_dir: Path): + # """Test edge cases in custom array extraction""" + # # Test with grid that has complex custom arrays that might cause extraction issues + # extended_grid = ExtendedGrid.empty() - # Add data that might cause issues during extraction - nodes = ExtendedNodeArray( - id=[1, 2], - u_rated=[10000, 10000], - u=[float("nan"), float("inf")], # Edge case values - ) - extended_grid.append(nodes) + # # Add data that might cause issues during extraction + # nodes = ExtendedNodeArray( + # id=[1, 2], + # u_rated=[10000, 10000], + # u=[float("nan"), float("inf")], # Edge case values + # ) + # extended_grid.append(nodes) - # Should handle edge case values gracefully - extensions = _extract_extensions_data(extended_grid) - assert "extended_columns" in extensions - assert "custom_arrays" in extensions + # # Should handle edge case values gracefully + # extensions = _extract_extensions_data(extended_grid) + # assert "extended_columns" in extensions + # assert "custom_arrays" in extensions - # Test saving and loading with these edge cases - json_path = temp_dir / "edge_cases.json" - save_grid_to_json(extended_grid, json_path, preserve_extensions=True) + # # Test saving and loading with these edge cases + # json_path = temp_dir / "edge_cases.json" + # save_grid_to_json(extended_grid, json_path, preserve_extensions=True) - # Should load without issues - loaded_grid = load_grid_from_json(json_path, target_grid_class=Grid) - assert loaded_grid.node.size == 2 + # # Should load without issues + # loaded_grid = load_grid_from_json(json_path, target_grid_class=Grid) + # assert loaded_grid.node.size == 2 def test_invalid_extension_data_recovery(self, temp_dir: Path): """Test recovery from invalid extension data""" @@ -265,7 +249,7 @@ def test_invalid_extension_data_recovery(self, temp_dir: Path): extended_grid.append(nodes) json_path = temp_dir / "test_recovery.json" - save_grid_to_json(extended_grid, json_path, preserve_extensions=True) + save_grid_to_json(extended_grid, json_path) # Corrupt extension data with open(json_path, "r") as f: From af873819928567c73e710cee6ed36658c8add85c Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Wed, 22 Oct 2025 15:29:39 +0200 Subject: [PATCH 09/19] feat: support additional values Signed-off-by: jaapschoutenalliander --- .../_core/utils/serialization.py | 29 ++++++++++++------- tests/unit/utils/test_serialization.py | 8 +++-- 2 files changed, 24 insertions(+), 13 deletions(-) diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py index c91b307..5834e09 100644 --- a/src/power_grid_model_ds/_core/utils/serialization.py +++ b/src/power_grid_model_ds/_core/utils/serialization.py @@ -16,25 +16,29 @@ logger = logging.getLogger(__name__) -def _restore_grid_arrays(grid, custom_arrays: Dict) -> None: - """Restore custom arrays to the grid.""" - for array_name, array_info in custom_arrays.items(): - if not hasattr(grid, array_name): +def _restore_grid_arrays(grid, input_data: Dict) -> None: + """Restore arrays to the grid.""" + for attr_name, attr_values in input_data.items(): + if not hasattr(grid, attr_name): + continue + + if not issubclass(getattr(grid, attr_name).__class__, FancyArray): + setattr(grid, attr_name, attr_values) continue try: - array_field = grid.find_array_field(getattr(grid, array_name).__class__) + array_field = grid.find_array_field(getattr(grid, attr_name).__class__) matched_columns = { - col: array_info["data"][col] for col in array_field.type().columns if col in array_info["data"] + 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, array_name, restored_array) + 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 custom array '{array_name}': {e}") + logger.warning(f"Failed to restore '{attr_name}': {e}") def save_grid_to_json( @@ -58,13 +62,16 @@ def save_grid_to_json( if field.name in ["graphs", "_id_counter"]: continue - array = getattr(grid, field.name) - if not isinstance(array, FancyArray) or array.size == 0: + field_value = getattr(grid, field.name) + if isinstance(field_value, (int, float, str, bool)): + serialized_data[field.name] = field_value + + if not isinstance(field_value, FancyArray) or field_value.size == 0: continue array_name = field.name serialized_data[array_name] = { - "data": {name: array[name].tolist() for name in array.dtype.names}, + "data": {name: field_value[name].tolist() for name in field_value.dtype.names}, } # Write to file diff --git a/tests/unit/utils/test_serialization.py b/tests/unit/utils/test_serialization.py index 9f353d8..9d8456b 100644 --- a/tests/unit/utils/test_serialization.py +++ b/tests/unit/utils/test_serialization.py @@ -46,8 +46,9 @@ class ExtendedGrid(Grid): node: ExtendedNodeArray line: ExtendedLineArray - # value_extension: float = 0.0 - # dict_extension: dict = dict() + value_extension: float = 0.0 + str_extension: str = "default" + complex_extension: list = None @pytest.fixture @@ -114,6 +115,9 @@ def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, tem # Verify core data assert loaded_grid.node.size == extended_grid.node.size assert loaded_grid.line.size == extended_grid.line.size + assert loaded_grid.value_extension == extended_grid.value_extension + assert loaded_grid.str_extension == extended_grid.str_extension + assert loaded_grid.complex_extension is None # Verify extended data np.testing.assert_array_equal(loaded_grid.node.u, extended_grid.node.u) From 89fa6eeb3e48dda7b1f76c40af611c4dbc0882d2 Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Wed, 22 Oct 2025 15:30:17 +0200 Subject: [PATCH 10/19] feat: support additional values Signed-off-by: jaapschoutenalliander --- src/power_grid_model_ds/_core/utils/serialization.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py index 5834e09..37ecf18 100644 --- a/src/power_grid_model_ds/_core/utils/serialization.py +++ b/src/power_grid_model_ds/_core/utils/serialization.py @@ -16,7 +16,7 @@ logger = logging.getLogger(__name__) -def _restore_grid_arrays(grid, input_data: Dict) -> None: +def _restore_grid_values(grid, input_data: Dict) -> None: """Restore arrays to the grid.""" for attr_name, attr_values in input_data.items(): if not hasattr(grid, attr_name): @@ -99,6 +99,6 @@ def load_grid_from_json(path: Path, target_grid_class=None): else: target_grid = target_grid_class.empty() - _restore_grid_arrays(target_grid, input_data) + _restore_grid_values(target_grid, input_data) return target_grid From e8385505ca544cf88f8bec14bc168dd6f23a0d98 Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Wed, 22 Oct 2025 15:32:01 +0200 Subject: [PATCH 11/19] feat: support additional values Signed-off-by: jaapschoutenalliander --- tests/unit/utils/test_serialization.py | 26 -------------------------- 1 file changed, 26 deletions(-) diff --git a/tests/unit/utils/test_serialization.py b/tests/unit/utils/test_serialization.py index 9d8456b..d5aa159 100644 --- a/tests/unit/utils/test_serialization.py +++ b/tests/unit/utils/test_serialization.py @@ -219,32 +219,6 @@ def test_empty_grid_handling(self, temp_dir: Path): loaded_json = load_grid_from_json(json_path, target_grid_class=Grid) assert loaded_json.node.size == 0 - # def test_custom_array_extraction_edge_cases(self, temp_dir: Path): - # """Test edge cases in custom array extraction""" - # # Test with grid that has complex custom arrays that might cause extraction issues - # extended_grid = ExtendedGrid.empty() - - # # Add data that might cause issues during extraction - # nodes = ExtendedNodeArray( - # id=[1, 2], - # u_rated=[10000, 10000], - # u=[float("nan"), float("inf")], # Edge case values - # ) - # extended_grid.append(nodes) - - # # Should handle edge case values gracefully - # extensions = _extract_extensions_data(extended_grid) - # assert "extended_columns" in extensions - # assert "custom_arrays" in extensions - - # # Test saving and loading with these edge cases - # json_path = temp_dir / "edge_cases.json" - # save_grid_to_json(extended_grid, json_path, preserve_extensions=True) - - # # Should load without issues - # loaded_grid = load_grid_from_json(json_path, target_grid_class=Grid) - # assert loaded_grid.node.size == 2 - def test_invalid_extension_data_recovery(self, temp_dir: Path): """Test recovery from invalid extension data""" # Create valid extended grid From 942e22b54574b320df903d9b46cd4c4034fb4e93 Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Wed, 22 Oct 2025 15:41:25 +0200 Subject: [PATCH 12/19] feat: remove obsolete test Signed-off-by: jaapschoutenalliander --- tests/unit/utils/test_serialization.py | 57 ++++++-------------------- 1 file changed, 13 insertions(+), 44 deletions(-) diff --git a/tests/unit/utils/test_serialization.py b/tests/unit/utils/test_serialization.py index d5aa159..33cb482 100644 --- a/tests/unit/utils/test_serialization.py +++ b/tests/unit/utils/test_serialization.py @@ -4,7 +4,6 @@ """Comprehensive unit tests for Grid serialization with power-grid-model compatibility.""" -import json from dataclasses import dataclass from pathlib import Path from tempfile import TemporaryDirectory @@ -123,6 +122,19 @@ def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, tem np.testing.assert_array_equal(loaded_grid.node.u, extended_grid.node.u) np.testing.assert_array_equal(loaded_grid.line.i_from, extended_grid.line.i_from) + def test_empty_grid_handling(self, temp_dir: Path): + """Test serialization of empty grids""" + empty_grid = Grid.empty() + + json_path = temp_dir / "empty.json" + + # Should handle empty grids + 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) + assert loaded_json.node.size == 0 + class TestCrossTypeCompatibility: """Test cross-type loading and compatibility""" @@ -201,46 +213,3 @@ class GridWithCustomArray(Grid): np.testing.assert_array_equal(loaded_grid.custom_metadata.id, [100, 200, 300]) np.testing.assert_array_almost_equal(loaded_grid.custom_metadata.metadata_value, [1.5, 2.5, 3.5]) np.testing.assert_array_equal(loaded_grid.custom_metadata.category, [1, 2, 1]) - - -class TestSpecialCases: - """Test special cases and edge scenarios""" - - def test_empty_grid_handling(self, temp_dir: Path): - """Test serialization of empty grids""" - empty_grid = Grid.empty() - - json_path = temp_dir / "empty.json" - - # Should handle empty grids - 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) - assert loaded_json.node.size == 0 - - def test_invalid_extension_data_recovery(self, temp_dir: Path): - """Test recovery from invalid extension data""" - # Create valid extended grid - extended_grid = ExtendedGrid.empty() - nodes = ExtendedNodeArray(id=[1, 2], u_rated=[10000, 10000], u=[9950, 9900]) - extended_grid.append(nodes) - - json_path = temp_dir / "test_recovery.json" - save_grid_to_json(extended_grid, json_path) - - # Corrupt extension data - with open(json_path, "r") as f: - data = json.load(f) - - # Add invalid extension data - if "pgm_ds_extensions" in data: - data["pgm_ds_extensions"]["extended_columns"]["node"]["u"] = [1, 2, 3, 4, 5] # Wrong size - data["pgm_ds_extensions"]["custom_arrays"]["fake"] = {"dtype": "invalid_dtype", "data": [[1, 2, 3]]} - - with open(json_path, "w") as f: - json.dump(data, f) - - # Should load core data despite extension errors - loaded_grid = load_grid_from_json(json_path, target_grid_class=Grid) - assert loaded_grid.node.size == 2 From ec8c44cfb7e6d937bd90e180202228df78acceb7 Mon Sep 17 00:00:00 2001 From: Jaap Schouten <58551444+jaapschoutenalliander@users.noreply.github.com> Date: Wed, 22 Oct 2025 18:14:52 +0200 Subject: [PATCH 13/19] chore: feedback Co-authored-by: Vincent Koppen <53343926+vincentkoppen@users.noreply.github.com> Signed-off-by: Jaap Schouten <58551444+jaapschoutenalliander@users.noreply.github.com> --- src/power_grid_model_ds/_core/utils/serialization.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py index 37ecf18..4e67308 100644 --- a/src/power_grid_model_ds/_core/utils/serialization.py +++ b/src/power_grid_model_ds/_core/utils/serialization.py @@ -69,8 +69,7 @@ def save_grid_to_json( if not isinstance(field_value, FancyArray) or field_value.size == 0: continue - array_name = field.name - serialized_data[array_name] = { + serialized_data[field.name] = { "data": {name: field_value[name].tolist() for name in field_value.dtype.names}, } From c588156eab0a5728dda3265491022a31e1ca58ae Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Wed, 22 Oct 2025 18:57:29 +0200 Subject: [PATCH 14/19] feat: pr feedback Signed-off-by: jaapschoutenalliander --- src/power_grid_model_ds/__init__.py | 6 -- .../_core/model/grids/base.py | 10 +++ .../_core/utils/serialization.py | 36 +++++---- tests/unit/utils/test_serialization.py | 77 ++++++++----------- 4 files changed, 66 insertions(+), 63 deletions(-) diff --git a/src/power_grid_model_ds/__init__.py b/src/power_grid_model_ds/__init__.py index 8638ffd..f1d7e6f 100644 --- a/src/power_grid_model_ds/__init__.py +++ b/src/power_grid_model_ds/__init__.py @@ -5,15 +5,9 @@ from power_grid_model_ds._core.load_flow import PowerGridModelInterface from power_grid_model_ds._core.model.graphs.container import GraphContainer from power_grid_model_ds._core.model.grids.base import Grid -from power_grid_model_ds._core.utils.serialization import ( - load_grid_from_json, - save_grid_to_json, -) __all__ = [ "Grid", "GraphContainer", "PowerGridModelInterface", - "save_grid_to_json", - "load_grid_from_json", ] diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py index ed586d8..9cc46d5 100644 --- a/src/power_grid_model_ds/_core/model/grids/base.py +++ b/src/power_grid_model_ds/_core/model/grids/base.py @@ -42,6 +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.zip import file2gzip Self = TypeVar("Self", bound="Grid") @@ -438,6 +439,15 @@ 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) -> Path: + """Serialize the grid to JSON format.""" + return _save_grid_to_json(grid=self, path=path) + + @classmethod + def from_json(cls: Type[Self], path: Path) -> Self: + """Deserialize the grid from JSON format.""" + 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""" set_is_feeder(grid=self) diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py index 4e67308..c324f3b 100644 --- a/src/power_grid_model_ds/_core/utils/serialization.py +++ b/src/power_grid_model_ds/_core/utils/serialization.py @@ -8,17 +8,25 @@ import json import logging from pathlib import Path -from typing import Dict, Optional +from typing import TYPE_CHECKING, Dict, Optional, Type, TypeVar from power_grid_model_ds._core.model.arrays.base.array import FancyArray -from power_grid_model_ds._core.model.grids.base import Grid + +if TYPE_CHECKING: + # Import only for type checking to avoid circular imports at runtime + from power_grid_model_ds._core.model.grids.base import Grid + + G = TypeVar("G", bound=Grid) +else: + # Runtime: don't import Grid to avoid circular import; keep unbound TypeVar + G = TypeVar("G") logger = logging.getLogger(__name__) -def _restore_grid_values(grid, input_data: Dict) -> None: +def _restore_grid_values(grid, json_data: Dict) -> None: """Restore arrays to the grid.""" - for attr_name, attr_values in input_data.items(): + for attr_name, attr_values in json_data.items(): if not hasattr(grid, attr_name): continue @@ -41,7 +49,7 @@ def _restore_grid_values(grid, input_data: Dict) -> None: logger.warning(f"Failed to restore '{attr_name}': {e}") -def save_grid_to_json( +def _save_grid_to_json( grid, path: Path, indent: Optional[int] = None, @@ -64,9 +72,13 @@ def save_grid_to_json( field_value = getattr(grid, field.name) if isinstance(field_value, (int, float, str, bool)): - serialized_data[field.name] = field_value + serialized_data[field.name] = field.type(field_value) + continue - if not isinstance(field_value, FancyArray) or field_value.size == 0: + 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: continue serialized_data[field.name] = { @@ -80,12 +92,12 @@ def save_grid_to_json( return path -def load_grid_from_json(path: Path, target_grid_class=None): +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: path: The file path to load from - target_grid_class: Optional Grid class to load into. If None, uses default Grid. + target_grid_class: Grid class to load into. Returns: Grid: The deserialized Grid object of the specified target class @@ -93,11 +105,7 @@ def load_grid_from_json(path: Path, target_grid_class=None): with open(path, "r", encoding="utf-8") as f: input_data = json.load(f) - if target_grid_class is None: - target_grid = Grid.empty() - else: - target_grid = target_grid_class.empty() - + target_grid = target_grid_class.empty() _restore_grid_values(target_grid, input_data) return target_grid diff --git a/tests/unit/utils/test_serialization.py b/tests/unit/utils/test_serialization.py index 33cb482..168b9b0 100644 --- a/tests/unit/utils/test_serialization.py +++ b/tests/unit/utils/test_serialization.py @@ -6,7 +6,6 @@ from dataclasses import dataclass from pathlib import Path -from tempfile import TemporaryDirectory import numpy as np import pytest @@ -15,11 +14,12 @@ 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 +from power_grid_model_ds.fancypy import array_equal class ExtendedNodeArray(BaseNodeArray): @@ -47,14 +47,6 @@ class ExtendedGrid(Grid): value_extension: float = 0.0 str_extension: str = "default" - complex_extension: list = None - - -@pytest.fixture -def temp_dir(): - """Temporary directory fixture""" - with TemporaryDirectory() as tmp_dir: - yield Path(tmp_dir) @pytest.fixture @@ -89,85 +81,84 @@ def extended_grid(): return grid -class TestSerializationFormats: +class TestSerializationRoundtrips: """Test serialization across different formats and configurations""" - def test_basic_serialization_roundtrip(self, basic_grid: Grid, temp_dir: Path): + def test_basic_serialization_roundtrip(self, basic_grid: Grid, tmp_path: Path): """Test basic serialization roundtrip for all formats""" - path = temp_dir / "test.json" - result_path = save_grid_to_json(basic_grid, path) + path = tmp_path / "test.json" + 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) - assert loaded_grid.node.size == basic_grid.node.size - assert loaded_grid.line.size == basic_grid.line.size + 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) - def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, temp_dir: Path): + def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, tmp_path: Path): """Test extended serialization preserving custom data""" - path = temp_dir / "extended.json" + 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 assert loaded_grid.line.size == extended_grid.line.size assert loaded_grid.value_extension == extended_grid.value_extension assert loaded_grid.str_extension == extended_grid.str_extension - assert loaded_grid.complex_extension is None # Verify extended data np.testing.assert_array_equal(loaded_grid.node.u, extended_grid.node.u) np.testing.assert_array_equal(loaded_grid.line.i_from, extended_grid.line.i_from) - def test_empty_grid_handling(self, temp_dir: Path): + def test_empty_grid_handling(self, tmp_path: Path): """Test serialization of empty grids""" empty_grid = Grid.empty() - json_path = temp_dir / "empty.json" + 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 class TestCrossTypeCompatibility: """Test cross-type loading and compatibility""" - def test_basic_to_extended_loading(self, basic_grid: Grid, temp_dir: Path): + def test_basic_to_extended_loading(self, basic_grid: Grid, tmp_path: Path): """Test loading basic grid into extended type""" - path = temp_dir / "basic.json" + 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 - assert loaded_grid.node.size == basic_grid.node.size - assert loaded_grid.line.size == basic_grid.line.size + array_equal(loaded_grid.node, basic_grid.node) + array_equal(loaded_grid.line, basic_grid.line) - def test_extended_to_basic_loading(self, extended_grid: ExtendedGrid, temp_dir: Path): + def test_extended_to_basic_loading(self, extended_grid: ExtendedGrid, tmp_path: Path): """Test loading extended grid into basic type""" - path = temp_dir / "extended.json" + 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 - assert loaded_grid.node.size == extended_grid.node.size - assert loaded_grid.line.size == extended_grid.line.size + array_equal(loaded_grid.node, extended_grid.node) + array_equal(loaded_grid.line, extended_grid.line) class TestExtensionHandling: """Test extension data handling and edge cases""" - def test_custom_array_serialization_roundtrip(self, temp_dir: Path): + def test_custom_array_serialization_roundtrip(self, tmp_path: Path): """Test serialization and loading of grids with custom arrays""" # Create a custom array type that properly extends FancyArray @@ -197,11 +188,11 @@ class GridWithCustomArray(Grid): grid.custom_metadata = custom_data # Test JSON serialization - json_path = temp_dir / "custom_array.json" - save_grid_to_json(grid, json_path) + json_path = tmp_path / "custom_array.json" + _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 From 8bcf0e71e8c8fd8b4ef2d64b49b8bf1b70a3cc7d Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Wed, 22 Oct 2025 19:01:55 +0200 Subject: [PATCH 15/19] feat: pr feedback json kwargs Signed-off-by: jaapschoutenalliander --- src/power_grid_model_ds/_core/model/grids/base.py | 13 ++++++++++--- .../_core/utils/serialization.py | 9 +++++---- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py index 9cc46d5..d0d11be 100644 --- a/src/power_grid_model_ds/_core/model/grids/base.py +++ b/src/power_grid_model_ds/_core/model/grids/base.py @@ -439,9 +439,16 @@ 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) -> Path: - """Serialize the grid to JSON format.""" - return _save_grid_to_json(grid=self, path=path) + def to_json(self, path: Path, **kwargs) -> Path: + """Serialize the grid to JSON format. + + Args: + path: Destination file path to write JSON to. + **kwargs: Additional keyword arguments forwarded to ``json.dump`` + Returns: + Path: The path where the file was saved. + """ + return _save_grid_to_json(grid=self, path=path, **kwargs) @classmethod def from_json(cls: Type[Self], path: Path) -> Self: diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py index c324f3b..f550dd3 100644 --- a/src/power_grid_model_ds/_core/utils/serialization.py +++ b/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, Optional, Type, TypeVar +from typing import TYPE_CHECKING, Dict, Type, TypeVar from power_grid_model_ds._core.model.arrays.base.array import FancyArray @@ -52,14 +52,15 @@ def _restore_grid_values(grid, json_data: Dict) -> None: def _save_grid_to_json( grid, path: Path, - indent: Optional[int] = None, + **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 - indent: JSON indentation (None for compact, positive int for indentation) + **kwargs: Keyword arguments forwarded to json.dump (for example, indent, sort_keys, + ensure_ascii, etc.). Returns: Path: The path where the file was saved """ @@ -87,7 +88,7 @@ def _save_grid_to_json( # Write to file with open(path, "w", encoding="utf-8") as f: - json.dump(serialized_data, f, indent=indent if indent and indent > 0 else None) + json.dump(serialized_data, f, **kwargs) return path From b1da441e576858042f522d536e6d6e2ca672a321 Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Wed, 22 Oct 2025 19:12:24 +0200 Subject: [PATCH 16/19] chore: type Signed-off-by: jaapschoutenalliander --- src/power_grid_model_ds/_core/utils/serialization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py index f550dd3..7594bc5 100644 --- a/src/power_grid_model_ds/_core/utils/serialization.py +++ b/src/power_grid_model_ds/_core/utils/serialization.py @@ -73,7 +73,7 @@ def _save_grid_to_json( field_value = getattr(grid, field.name) if isinstance(field_value, (int, float, str, bool)): - serialized_data[field.name] = field.type(field_value) + serialized_data[field.name] = field.type(field_value) # type: ignore continue if not isinstance(field_value, FancyArray): From c35e081abd60112c1adead55d24acf82e7926bcd Mon Sep 17 00:00:00 2001 From: jaapschoutenalliander Date: Wed, 22 Oct 2025 19:17:54 +0200 Subject: [PATCH 17/19] chore: type Signed-off-by: jaapschoutenalliander --- src/power_grid_model_ds/_core/utils/serialization.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py index 7594bc5..2e4ea95 100644 --- a/src/power_grid_model_ds/_core/utils/serialization.py +++ b/src/power_grid_model_ds/_core/utils/serialization.py @@ -31,7 +31,9 @@ def _restore_grid_values(grid, json_data: Dict) -> None: continue if not issubclass(getattr(grid, attr_name).__class__, FancyArray): - setattr(grid, attr_name, attr_values) + expected_type = grid.__dataclass_fields__[attr_name].type + cast_value = expected_type(attr_values) + setattr(grid, attr_name, cast_value) continue try: @@ -73,7 +75,7 @@ def _save_grid_to_json( field_value = getattr(grid, field.name) if isinstance(field_value, (int, float, str, bool)): - serialized_data[field.name] = field.type(field_value) # type: ignore + serialized_data[field.name] = field_value continue if not isinstance(field_value, FancyArray): From 257360d14f2d72b713182f632f9c9ed88a50a552 Mon Sep 17 00:00:00 2001 From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com> Date: Wed, 12 Nov 2025 16:41:57 +0100 Subject: [PATCH 18/19] Apply some changes from review Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com> --- .../_core/model/grids/base.py | 10 +-- .../_core/utils/serialization.py | 88 +++++++++---------- tests/integration/visualizer_tests.py | 6 ++ tests/unit/utils/test_serialization.py | 28 +++--- 4 files changed, 66 insertions(+), 66 deletions(-) diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py index 33d7616..d19b3bf 100644 --- a/src/power_grid_model_ds/_core/model/grids/base.py +++ b/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""" diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py index 2e4ea95..cf67669 100644 --- a/src/power_grid_model_ds/_core/utils/serialization.py +++ b/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)) diff --git a/tests/integration/visualizer_tests.py b/tests/integration/visualizer_tests.py index 2cb437a..9dbabdf 100644 --- a/tests/integration/visualizer_tests.py +++ b/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() diff --git a/tests/unit/utils/test_serialization.py b/tests/unit/utils/test_serialization.py index 168b9b0..c117ad8 100644 --- a/tests/unit/utils/test_serialization.py +++ b/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 From deb0c57b7430cc803c4042a92e33831d3c8b5575 Mon Sep 17 00:00:00 2001 From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com> Date: Wed, 12 Nov 2025 22:37:33 +0100 Subject: [PATCH 19/19] Apply suggestion from @Thijss Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com> --- src/power_grid_model_ds/_core/model/grids/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py index d19b3bf..6f4bf97 100644 --- a/src/power_grid_model_ds/_core/model/grids/base.py +++ b/src/power_grid_model_ds/_core/model/grids/base.py @@ -363,7 +363,7 @@ def get_downstream_nodes(self, node_id: int, inclusive: bool = False): def cache(self, cache_dir: Path, cache_name: str, compress: bool = True): """Cache Grid to a folder using pickle format. - Note: Consider using save_to_json() for better + Note: Consider using serialize() for better interoperability and standardized format. Args: