diff --git a/docs/changes/newsfragments/6599.new b/docs/changes/newsfragments/6599.new new file mode 100644 index 00000000000..5bff7f1693a --- /dev/null +++ b/docs/changes/newsfragments/6599.new @@ -0,0 +1 @@ +Added a new feature to find links in a parameter chain by parameter type diff --git a/src/qcodes/extensions/__init__.py b/src/qcodes/extensions/__init__.py index a18aba947c0..9cb766d3cbe 100644 --- a/src/qcodes/extensions/__init__.py +++ b/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", diff --git a/src/qcodes/extensions/infer.py b/src/qcodes/extensions/infer.py index e3361fa4fff..df2e5eb7ff1 100644 --- a/src/qcodes/extensions/infer.py +++ b/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] diff --git a/tests/extensions/test_infer.py b/tests/extensions/test_infer.py index d0453b96e12..13920ff68b4 100644 --- a/tests/extensions/test_infer.py +++ b/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)