Merge pull request #6599 from samantha-ho/samanthaho/parameter_chain_…

…link_by_type

Add methods to find parameter chain links by type

docs/changes/newsfragments/6599.new

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

src/qcodes/extensions/__init__.py

@@ -8,7 +8,11 @@
 from .infer import (
     InferAttrs,
     InferError,
+    get_chain_links_of_type,
+    get_instrument_from_param,
+    get_parameter_chain,
     get_root_parameter,
+    get_sole_chain_link_of_type,
     infer_channel,
     infer_instrument,
     infer_instrument_module,
@@ -21,7 +25,11 @@
     "DriverTestCase",
     "InferAttrs",
     "InferError",
+    "get_chain_links_of_type",
+    "get_instrument_from_param",
+    "get_parameter_chain",
     "get_root_parameter",
+    "get_sole_chain_link_of_type",
     "infer_channel",
     "infer_instrument",
     "infer_instrument_module",

src/qcodes/extensions/infer.py

@@ -1,16 +1,18 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, ClassVar
+from typing import TYPE_CHECKING, ClassVar, TypeVar, cast
 
 from qcodes.instrument import Instrument, InstrumentBase, InstrumentModule
-from qcodes.parameters import DelegateParameter, Parameter
+from qcodes.parameters import DelegateParameter, Parameter, ParameterBase
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
 DOES_NOT_EXIST = "Does not exist"
 
+C = TypeVar("C", bound=ParameterBase)
+
 
 class InferError(AttributeError): ...
 
@@ -220,3 +222,34 @@ def _merge_user_and_class_attrs(
         return set.union(set((alt_source_attrs,)), set(InferAttrs.known_attrs()))
     else:
         return set.union(set(alt_source_attrs), set(InferAttrs.known_attrs()))
+
+
+def get_chain_links_of_type(
+    link_param_type: type[C] | tuple[type[C], ...], parameter: Parameter
+) -> tuple[C, ...]:
+    """Gets all parameters in a chain of linked parameters that match a given type"""
+    chain_links: list[C] = [
+        cast(C, param)
+        for param in get_parameter_chain(parameter)
+        if isinstance(param, link_param_type)
+    ]
+    return tuple(chain_links)
+
+
+def get_sole_chain_link_of_type(
+    link_param_type: type[C] | tuple[type[C], ...], parameter: Parameter
+) -> C:
+    """Gets the one parameter in a chain of linked parameters that matches a given type"""
+
+    chain_links = get_chain_links_of_type(
+        link_param_type=link_param_type, parameter=parameter
+    )
+    if len(chain_links) != 1:
+        if isinstance(link_param_type, type):
+            error_msg_1 = f"Expected only a single chain link of type {link_param_type.__name__} but found {len(chain_links)}: \n"
+        elif isinstance(link_param_type, tuple):
+            type_strs = [link_type.__name__ for link_type in link_param_type]
+            error_msg_1 = f"Expected only a single chain link of types {type_strs} but found {len(chain_links)}: \n"
+
+        raise ValueError(error_msg_1 + f"{[link.name for link in chain_links]}")
+    return chain_links[0]

tests/extensions/test_infer.py

@@ -9,8 +9,10 @@
     InferAttrs,
     InferError,
     _merge_user_and_class_attrs,
+    get_chain_links_of_type,
     get_parameter_chain,
     get_root_parameter,
+    get_sole_chain_link_of_type,
     infer_channel,
     infer_instrument,
 )
@@ -280,3 +282,35 @@ def test_get_root_parameter_with_loops(good_inst_delegates):
     with pytest.raises(InferError) as exc_info:
         get_root_parameter(good_inst_del_2, "linked_parameter")
     assert "generated a loop of linking parameters" in str(exc_info.value)
+
+
+def test_chain_links_of_type():
+    root_param = ManualParameter("root", initial_value=0)
+    user_link = UserLinkingParameter("user_link", linked_parameter=root_param)
+    delegate_link = DelegateParameter("delegate_link", source=user_link)
+    user_link_2 = UserLinkingParameter("user_link_2", linked_parameter=delegate_link)
+
+    InferAttrs.add("linked_parameter")
+    user_links = get_chain_links_of_type(UserLinkingParameter, user_link_2)
+    assert set(user_links) == set([user_link, user_link_2])
+
+
+def test_sole_chain_link_of_type():
+    root_param = ManualParameter("root", initial_value=0)
+    user_link = UserLinkingParameter("user_link", linked_parameter=root_param)
+    delegate_link = DelegateParameter("delegate_link", source=user_link)
+    user_link_2 = UserLinkingParameter("user_link_2", linked_parameter=delegate_link)
+
+    InferAttrs.add("linked_parameter")
+    sole_user_link = get_sole_chain_link_of_type(UserLinkingParameter, delegate_link)
+    assert sole_user_link == user_link
+
+    with pytest.raises(ValueError) as exc_info:
+        _ = get_sole_chain_link_of_type(UserLinkingParameter, user_link_2)
+    assert "Expected only a single chain link of type" in str(exc_info.value)
+
+    with pytest.raises(ValueError) as exc_info:
+        _ = get_sole_chain_link_of_type(
+            (UserLinkingParameter, DelegateParameter), user_link_2
+        )
+    assert "Expected only a single chain link of types" in str(exc_info.value)