From b7eade81b87e40538d0fda8e9060cb00ff9f20e1 Mon Sep 17 00:00:00 2001
From: Serwan Asaad <serwan.asaad@quantum-machines.co>
Date: Thu, 4 Jul 2024 04:47:01 +0800
Subject: [PATCH 1/3] add converter for int keys

---
 quam/serialisation/json.py                    | 16 +++++++++++-
 .../serialisation/test_json_serialisation.py  | 26 +++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

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

From abff173cd9658cba0a6a9cd0a95791b1b53ef394 Mon Sep 17 00:00:00 2001
From: Serwan Asaad <serwan.asaad@quantum-machines.co>
Date: Thu, 4 Jul 2024 05:04:04 +0800
Subject: [PATCH 2/3] Update changelog

---
 CHANGELOG.md | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 99f46e16..dc3496ae 100644
--- a/CHANGELOG.md
+++ b/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`

From ee64c8c7d5f3d72516924f82758cb723a70a8aef Mon Sep 17 00:00:00 2001
From: Serwan Asaad <serwan.asaad@quantum-machines.co>
Date: Thu, 4 Jul 2024 06:29:39 +0800
Subject: [PATCH 3/3] test fix

---
 tests/serialisation/test_json_serialisation.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/serialisation/test_json_serialisation.py b/tests/serialisation/test_json_serialisation.py
index b2ff446a..167f4c4d 100644
--- a/tests/serialisation/test_json_serialisation.py
+++ b/tests/serialisation/test_json_serialisation.py
@@ -128,8 +128,8 @@ def test_serialise_int_dict_keys(tmp_path):
     assert d == {
         "a": 1,
         "d": {
-            "1": "a",
-            "2": "b",
+            1: "a",
+            2: "b",
         },
         "__class__": "test_json_serialisation.QuAMWithIntDict",
     }