Merge pull request #6616 from picarro-yren/yure/get_intrument

Add methods to find instruments from parameter chain

docs/changes/newsfragments/6616.new

@@ -0,0 +1 @@
+Added new methods to find instruments in a parameter chain by parameter type

src/qcodes/extensions/__init__.py

@@ -11,8 +11,10 @@
     get_chain_links_of_type,
     get_instrument_from_param,
     get_parameter_chain,
+    get_parent_instruments_from_chain_of_type,
     get_root_parameter,
     get_sole_chain_link_of_type,
+    get_sole_parent_instrument_from_chain_of_type,
     infer_channel,
     infer_instrument,
     infer_instrument_module,
@@ -30,6 +32,8 @@
     "get_parameter_chain",
     "get_root_parameter",
     "get_sole_chain_link_of_type",
+    "get_parent_instruments_from_chain_of_type",
+    "get_sole_parent_instrument_from_chain_of_type",
     "infer_channel",
     "infer_instrument",
     "infer_instrument_module",

src/qcodes/extensions/infer.py

@@ -12,6 +12,7 @@
 DOES_NOT_EXIST = "Does not exist"
 
 C = TypeVar("C", bound=ParameterBase)
+TInstrument = TypeVar("TInstrument", bound=InstrumentBase)
 
 
 class InferError(AttributeError): ...
@@ -253,3 +254,38 @@ def get_sole_chain_link_of_type(
 
         raise ValueError(error_msg_1 + f"{[link.name for link in chain_links]}")
     return chain_links[0]
+
+
+def get_parent_instruments_from_chain_of_type(
+    instrument_type: type[TInstrument] | tuple[type[TInstrument], ...],
+    parameter: Parameter,
+) -> tuple[TInstrument, ...]:
+    """Gets all parent instruments in a chain of linked parameters that match a given type"""
+
+    param_chain = get_parameter_chain(parameter)
+    return tuple(
+        [
+            cast(TInstrument, param.instrument)
+            for param in param_chain
+            if isinstance(param.instrument, instrument_type)
+        ]
+    )
+
+
+def get_sole_parent_instrument_from_chain_of_type(
+    instrument_type: type[TInstrument] | tuple[type[TInstrument], ...],
+    parameter: Parameter,
+) -> TInstrument:
+    """Gets the one parent instruments in a chain of linked parameters that match a given type"""
+    instruments = get_parent_instruments_from_chain_of_type(
+        instrument_type=instrument_type, parameter=parameter
+    )
+    if len(instruments) != 1:
+        if isinstance(instrument_type, type):
+            error_msg_1 = f"Expected only a single instrument of type {instrument_type.__name__} but found {len(instruments)}: \n"
+        elif isinstance(instrument_type, tuple):
+            type_strs = [instr_type.__name__ for instr_type in instrument_type]
+            error_msg_1 = f"Expected only a single instrument of types {type_strs} but found {len(instruments)}: \n"
+
+        raise ValueError(f"{error_msg_1} {[instr.name for instr in instruments]}")
+    return instruments[0]

tests/extensions/test_infer.py

@@ -11,8 +11,10 @@
     _merge_user_and_class_attrs,
     get_chain_links_of_type,
     get_parameter_chain,
+    get_parent_instruments_from_chain_of_type,
     get_root_parameter,
     get_sole_chain_link_of_type,
+    get_sole_parent_instrument_from_chain_of_type,
     infer_channel,
     infer_instrument,
 )
@@ -39,6 +41,11 @@ def __init__(self, name: str):
         self.module = DummyModule(name="module", parent=self)
 
 
+class DummyInstrument2(DummyInstrument):
+    def __init__(self, name: str):
+        super().__init__(name=name)
+
+
 class DummyDelegateInstrument(InstrumentBase):
     def __init__(self, name: str):
         super().__init__(name=name)
@@ -77,6 +84,16 @@ def make_instrument_fixture():
         inst.close()
 
 
+@pytest.fixture(name="instrument_fixture2")
+def make_instrument_fixture2():
+    inst = DummyInstrument2("dummy_instrument2")
+    InferAttrs.clear()
+    try:
+        yield inst
+    finally:
+        inst.close()
+
+
 @pytest.fixture(name="good_inst_delegates")
 def make_good_delegate_parameters(instrument_fixture):
     inst = instrument_fixture
@@ -90,6 +107,24 @@ def make_good_delegate_parameters(instrument_fixture):
     return good_inst_del_1, good_inst_del_2, good_inst_del_3
 
 
+@pytest.fixture(name="multi_inst_chain")
+def make_multi_instrument_chain(instrument_fixture, instrument_fixture2):
+    inst = instrument_fixture
+    inst2 = instrument_fixture2
+    good_inst_del_1 = DelegateParameter(
+        "good_inst_del_1", source=inst.good_inst_parameter
+    )
+    good_inst_del_2 = DelegateParameter(
+        "good_inst_del_2",
+        source=good_inst_del_1,
+        instrument=inst2,
+    )
+    good_inst_del_3 = UserLinkingParameter(
+        "good_inst_del_3", linked_parameter=good_inst_del_2
+    )
+    return good_inst_del_1, good_inst_del_2, good_inst_del_3
+
+
 def test_get_root_parameter_valid(instrument_fixture, good_inst_delegates):
     inst = instrument_fixture
     good_inst_del_1, good_inst_del_2, good_inst_del_3 = good_inst_delegates
@@ -314,3 +349,40 @@ def test_sole_chain_link_of_type():
             (UserLinkingParameter, DelegateParameter), user_link_2
         )
     assert "Expected only a single chain link of types" in str(exc_info.value)
+
+
+def test_get_instrument_from_chain(
+    instrument_fixture, instrument_fixture2, multi_inst_chain
+):
+    inst = instrument_fixture
+    inst2 = instrument_fixture2
+    good_inst_del_1, good_inst_del_2, good_inst_del_3 = multi_inst_chain
+
+    InferAttrs.add("linked_parameter")
+    instruments = get_parent_instruments_from_chain_of_type(
+        DummyInstrument, good_inst_del_3
+    )
+    assert set(instruments) == set([inst, inst2])
+
+
+def test_get_sole_instrument_from_chain(instrument_fixture2, multi_inst_chain):
+    inst2 = instrument_fixture2
+    good_inst_del_1, good_inst_del_2, good_inst_del_3 = multi_inst_chain
+
+    InferAttrs.add("linked_parameter")
+    sole_instrument = get_sole_parent_instrument_from_chain_of_type(
+        DummyInstrument2, good_inst_del_3
+    )
+    assert sole_instrument == inst2
+
+    with pytest.raises(ValueError) as exc_info:
+        _ = get_sole_parent_instrument_from_chain_of_type(
+            DummyInstrument, good_inst_del_3
+        )
+    assert "Expected only a single instrument of type" in str(exc_info.value)
+
+    with pytest.raises(ValueError) as exc_info:
+        _ = get_sole_parent_instrument_from_chain_of_type(
+            (DummyInstrument, DummyInstrument2), good_inst_del_3
+        )
+    assert "Expected only a single instrument of types" in str(exc_info.value)