From 487cfd77979d9473082b9abd007def5a37816961 Mon Sep 17 00:00:00 2001 From: Serwan Asaad Date: Fri, 8 Mar 2024 08:26:08 +0100 Subject: [PATCH] Support int keys for QuamDict --- quam/core/quam_classes.py | 2 +- quam/utils/string_reference.py | 8 +++++--- .../quam_base/referencing/test_referencing_dict.py | 14 ++++++++++++++ tests/quam_base/test_quam_dict.py | 10 ++++++++++ 4 files changed, 30 insertions(+), 4 deletions(-) diff --git a/quam/core/quam_classes.py b/quam/core/quam_classes.py index dbfc6910..7e2455c3 100644 --- a/quam/core/quam_classes.py +++ b/quam/core/quam_classes.py @@ -73,7 +73,7 @@ def convert_dict_and_list(value, cls_or_obj=None, attr=None): """Convert a dict or list to a QuamDict or QuamList if possible.""" if isinstance(value, dict): value_annotation = _get_value_annotation(cls_or_obj=cls_or_obj, attr=attr) - return QuamDict(**value, value_annotation=value_annotation) + return QuamDict(value, value_annotation=value_annotation) elif type(value) == list: value_annotation = _get_value_annotation(cls_or_obj=cls_or_obj, attr=attr) return QuamList(value, value_annotation=value_annotation) diff --git a/quam/utils/string_reference.py b/quam/utils/string_reference.py index 9240598e..6a7f65b5 100644 --- a/quam/utils/string_reference.py +++ b/quam/utils/string_reference.py @@ -79,10 +79,12 @@ def get_relative_reference_value(obj, string: str) -> Any: except KeyError as e: raise AttributeError(f"Object {obj} has no attribute {next_attr}") from e elif isinstance(obj, (dict, UserDict)): - try: + if next_attr in obj: obj_attr = obj[next_attr] - except KeyError as e: - raise AttributeError(f"Object {obj} has no attribute {next_attr}") from e + elif next_attr.isdigit() and int(next_attr) in obj: + obj_attr = obj[int(next_attr)] + else: + raise AttributeError(f"Object {obj} has no attribute {next_attr}") else: obj_attr = getattr(obj, next_attr) diff --git a/tests/quam_base/referencing/test_referencing_dict.py b/tests/quam_base/referencing/test_referencing_dict.py index 899006ae..1f9a2d5d 100644 --- a/tests/quam_base/referencing/test_referencing_dict.py +++ b/tests/quam_base/referencing/test_referencing_dict.py @@ -46,3 +46,17 @@ def test_referencing_to_dict(): assert quam_root._get_referenced_value("#/quam_dict/a") == 44 assert quam_root.quam_dict._get_referenced_value("#/quam_dict/a") == 44 assert quam_root.quam_dict._get_referenced_value("#./a") == 44 + + +def test_referencing_dict_int_keys(): + quam_root = BareQuamRoot() + quam_root.quam_dict = {"1": 1, 2: 2} + + assert list(quam_root.quam_dict.keys()) == list(quam_root.quam_dict) + assert list(quam_root.quam_dict) == ["1", 2] + + assert quam_root._get_referenced_value("#/quam_dict/1") == 1 + assert quam_root.quam_dict._get_referenced_value("#/quam_dict/1") == 1 + + assert quam_root._get_referenced_value("#/quam_dict/2") == 2 + assert quam_root.quam_dict._get_referenced_value("#/quam_dict/2") == 2 diff --git a/tests/quam_base/test_quam_dict.py b/tests/quam_base/test_quam_dict.py index 78e2c4ef..ea0980df 100644 --- a/tests/quam_base/test_quam_dict.py +++ b/tests/quam_base/test_quam_dict.py @@ -201,3 +201,13 @@ def test_dict_unreferenced_value(): assert d.val1 == 42 assert d.val2 == 42 assert d.get_unreferenced_value("val1") == "#./val2" + + +def test_quam_dict_int_keys(): + quam_dict = QuamDict({1: 2}) + assert quam_dict.data == {1: 2} + assert quam_dict[1] == 2 + quam_dict.pop(1) + assert quam_dict.data == {} + with pytest.raises(KeyError): + quam_dict[1]