Merge branch 'main' into fix/instantiate-union-pipe

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]

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

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",
+    }