Merge pull request #48 from qua-platform/fix/json-int-keys

add converter for int keys

CHANGELOG.md

@@ -6,6 +6,9 @@
 ### Changed
 - Allow `QuamBase.get_reference(attr)` to return a reference of one of its attributes
 
+### Fixed
+- Allow int keys to be serialised / loaded in QuAM using JSONSerialiser
+
 ## [0.3.3]
 ### Added
 - Added the following parameters to `IQChannel`: `RF_frequency`, `LO_frequency`, `intermediate_frequency`

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