diff --git a/CHANGELOG.md b/CHANGELOG.md index c873f676..12186a34 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ ### Fixed - Fix quam object instantiation error when a parameter type uses pipe operator +- Allow int keys to be serialised / loaded in QuAM using JSONSerialiser ## [0.3.3] diff --git a/quam/serialisation/json.py b/quam/serialisation/json.py index 3d7b9306..e2d56655 100644 --- a/quam/serialisation/json.py +++ b/quam/serialisation/json.py @@ -166,7 +166,7 @@ def load( metadata["default_filename"] = path.name with open(path, "r") as f: - contents = json.load(f) + contents = json.load(f, object_hook=convert_int_keys) elif path.is_dir(): metadata["default_foldername"] = str(path) for file in path.iterdir(): @@ -183,3 +183,17 @@ def load( metadata["content_mapping"][file.name] = list(file_contents.keys()) return contents, metadata + + +def convert_int_keys(obj): + """Convert dictionary keys to integers if possible.""" + if not isinstance(obj, dict): + return obj + + new_obj = {} + for key, value in obj.items(): + if key.isdigit(): + key = int(key) + new_obj[key] = value + + return new_obj diff --git a/tests/serialisation/test_json_serialisation.py b/tests/serialisation/test_json_serialisation.py index b7695e5f..167f4c4d 100644 --- a/tests/serialisation/test_json_serialisation.py +++ b/tests/serialisation/test_json_serialisation.py @@ -1,4 +1,5 @@ import json +from typing import Dict import pytest from quam.serialisation import JSONSerialiser @@ -107,3 +108,28 @@ def test_component_mamping_ignore(tmp_path): "a": 4, } } + + +@quam_dataclass +class QuAMWithIntDict(QuamRoot): + a: int + d: Dict[int, str] + + +def test_serialise_int_dict_keys(tmp_path): + quam_root = QuAMWithIntDict(a=1, d={1: "a", 2: "b"}) + + serialiser = JSONSerialiser() + path = tmp_path / "quam_root.json" + serialiser.save(quam_root, path) + + d, _ = serialiser.load(path) + + assert d == { + "a": 1, + "d": { + 1: "a", + 2: "b", + }, + "__class__": "test_json_serialisation.QuAMWithIntDict", + }