diff --git a/docs/changes/newsfragments/5998.new b/docs/changes/newsfragments/5998.new new file mode 100644 index 00000000000..b68ee0a99c0 --- /dev/null +++ b/docs/changes/newsfragments/5998.new @@ -0,0 +1,4 @@ +Add methods to recursively search a chain of DelegateParameters and return either all the parameters in the chain or the 'root' parameter +These methods may also be used with custom Parameters which link to other parameters via different attribute names + +Also add infer_channel and infer_instrument methods to find the InstrumentModule or Instrument of the root parameter diff --git a/src/qcodes/extensions/__init__.py b/src/qcodes/extensions/__init__.py index 9fee703eeac..a18aba947c0 100644 --- a/src/qcodes/extensions/__init__.py +++ b/src/qcodes/extensions/__init__.py @@ -2,12 +2,27 @@ The extensions module contains smaller modules that extend the functionality of QCoDeS. These modules may import from all of QCoDeS but do not themselves get imported into QCoDeS. """ + from ._driver_test_case import DriverTestCase from ._log_export_info import log_dataset_export_info +from .infer import ( + InferAttrs, + InferError, + get_root_parameter, + infer_channel, + infer_instrument, + infer_instrument_module, +) from .installation import register_station_schema_with_vscode __all__ = [ "register_station_schema_with_vscode", "log_dataset_export_info", "DriverTestCase", + "InferAttrs", + "InferError", + "get_root_parameter", + "infer_channel", + "infer_instrument", + "infer_instrument_module", ] diff --git a/src/qcodes/extensions/infer.py b/src/qcodes/extensions/infer.py new file mode 100644 index 00000000000..1937e20da6a --- /dev/null +++ b/src/qcodes/extensions/infer.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING, ClassVar + +from qcodes.instrument import Instrument, InstrumentBase, InstrumentModule +from qcodes.parameters import DelegateParameter, Parameter + +if TYPE_CHECKING: + from collections.abc import Iterable + +DOES_NOT_EXIST = "Does not exist" + + +class InferError(AttributeError): ... + + +class InferAttrs: + """Holds a global set of attribute name that will be inferred""" + + _known_attrs: ClassVar[set[str]] = set() + + @classmethod + def add(cls, attrs: str | Iterable[str]) -> None: + if isinstance(attrs, str): + attrs = (attrs,) + cls._known_attrs.update(set(attrs)) + + @classmethod + def known_attrs(cls) -> tuple[str, ...]: + return tuple(cls._known_attrs) + + @classmethod + def discard(cls, attr: str) -> None: + cls._known_attrs.discard(attr) + + @classmethod + def clear(cls) -> None: + cls._known_attrs = set() + + +def get_root_parameter( + param: Parameter, + alt_source_attrs: Sequence[str] | None = None, +) -> Parameter: + """ + Return the root parameter in a chain of DelegateParameters or other linking Parameters + + This method calls get_parameter_chain and then checks for various error conditions + Args: + param: The DelegateParameter or other linking parameter to find the root parameter from + alt_source_attrs: The attribute names for custom linking parameters + + Raises: + InferError: If the linking parameters do not end with a non-linking parameter + InferError: If the chain of linking parameters loops on itself + """ + + parameter_chain = get_parameter_chain(param, alt_source_attrs) + root_param = parameter_chain[-1] + + if root_param is parameter_chain[0] and len(parameter_chain) > 1: + raise InferError(f"{param} generated a loop of linking parameters") + if isinstance(root_param, DelegateParameter): + raise InferError(f"Parameter {param} is not attached to a source") + + alt_source_attrs_set = _merge_user_and_class_attrs(alt_source_attrs) + for alt_source_attr in alt_source_attrs_set: + alt_source = getattr(param, alt_source_attr, DOES_NOT_EXIST) + if alt_source is None: + raise InferError( + f"Parameter {param} is not attached to a source on attribute {alt_source_attr}" + ) + return root_param + + +def infer_instrument( + param: Parameter, + alt_source_attrs: Sequence[str] | None = None, +) -> InstrumentBase: + """ + Find the instrument that owns a parameter or delegate parameter. + + Args: + param: The DelegateParameter or other linking parameter to find the instrument from + alt_source_attrs: The attribute names for custom linking parameters + + Raises: + InferError: If the linking parameters do not end with a non-linking parameter + InferError: If the instrument of the root parameter is None + InferError: If the instrument of the root parameter is not an instance of Instrument + """ + root_param = get_root_parameter(param, alt_source_attrs=alt_source_attrs) + instrument = get_instrument_from_param(root_param) + if isinstance(instrument, InstrumentModule): + return instrument.root_instrument + elif isinstance(instrument, Instrument): + return instrument + + raise InferError(f"Could not determine source instrument for parameter {param}") + + +def infer_instrument_module( + param: Parameter, + alt_source_attrs: Sequence[str] | None = None, +) -> InstrumentModule: + """ + Find the instrument module that owns a parameter or delegate parameter + + Args: + param: The DelegateParameter or other linking parameter to find the instrument module from + alt_source_attrs: The attribute names for custom linking parameters + + Raises: + InferError: If the linking parameters do not end with a non-linking parameter + InferError: If the instrument module of the root parameter is None + InferError: If the instrument module of the root parameter is not an instance of InstrumentModule + """ + root_param = get_root_parameter(param, alt_source_attrs=alt_source_attrs) + channel = get_instrument_from_param(root_param) + if isinstance(channel, InstrumentModule): + return channel + raise InferError( + f"Could not determine a root instrument channel for parameter {param}" + ) + + +def infer_channel( + param: Parameter, + alt_source_attrs: Sequence[str] | None = None, +) -> InstrumentModule: + """An alias for infer_instrument_module""" + return infer_instrument_module(param, alt_source_attrs) + + +def get_instrument_from_param( + param: Parameter, +) -> InstrumentBase: + """ + Return the instrument attribute from a parameter + + Args: + param: The parameter to get the instrument module from + + Raises: + InferError: If the parameter does not have an instrument + """ + if param.instrument is not None: + return param.instrument + raise InferError(f"Parameter {param} has no instrument") + + +def get_parameter_chain( + param_chain: Parameter | Sequence[Parameter], + alt_source_attrs: str | Sequence[str] | None = None, +) -> tuple[Parameter, ...]: + """ + Return the chain of DelegateParameters or other linking Parameters + + This method traverses singly-linked parameters and returns the resulting chain + If the parameters loop, then the first and last linking parameters in the chain + will be identical. Otherwise, the chain starts with the initial argument passed + and ends when the chain terminates in either a non-linking parameter or a + linking parameter that links to None + + The search prioritizes the `source` attribute of DelegateParameters first, and + then looks for other linking attributes in undetermined order. + + Args: + param_chain: The initial linking parameter or a List linking parameters + from which to return the chain + alt_source_attrs: The attribute names for custom linking parameters + """ + + alt_source_attrs_set = _merge_user_and_class_attrs(alt_source_attrs) + + if not isinstance(param_chain, Sequence): + param_chain = (param_chain,) + + param = param_chain[-1] + mutable_param_chain = list(param_chain) + if isinstance(param, DelegateParameter): + if param.source is None: + return tuple(param_chain) + mutable_param_chain.append(param.source) + if param.source in param_chain: # There is a loop in the links + return tuple(mutable_param_chain) + return get_parameter_chain( + mutable_param_chain, + alt_source_attrs=alt_source_attrs, + ) + + for alt_source_attr in alt_source_attrs_set: + alt_source = getattr(param, alt_source_attr, DOES_NOT_EXIST) + if alt_source is None: # Valid linking attribute, but no link parameter + return tuple(param_chain) + elif isinstance(alt_source, Parameter): + mutable_param_chain.append(alt_source) + if alt_source in param_chain: # There is a loop in the links + return tuple(mutable_param_chain) + return get_parameter_chain( + mutable_param_chain, + alt_source_attrs=alt_source_attrs, + ) + return tuple(param_chain) + + +def _merge_user_and_class_attrs( + alt_source_attrs: str | Sequence[str] | None = None, +) -> Iterable[str]: + """Merges user-supplied linking attributes with attributes from InferAttrs""" + if alt_source_attrs is None: + return InferAttrs.known_attrs() + elif isinstance(alt_source_attrs, str): + return set.union(set((alt_source_attrs,)), set(InferAttrs.known_attrs())) + else: + return set.union(set(alt_source_attrs), set(InferAttrs.known_attrs())) diff --git a/tests/extensions/test_infer.py b/tests/extensions/test_infer.py new file mode 100644 index 00000000000..d0453b96e12 --- /dev/null +++ b/tests/extensions/test_infer.py @@ -0,0 +1,282 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import pytest + +from qcodes.extensions.infer import ( + InferAttrs, + InferError, + _merge_user_and_class_attrs, + get_parameter_chain, + get_root_parameter, + infer_channel, + infer_instrument, +) +from qcodes.instrument import Instrument, InstrumentBase, InstrumentModule +from qcodes.parameters import DelegateParameter, ManualParameter, Parameter + + +class DummyModule(InstrumentModule): + def __init__(self, name: str, parent: Instrument): + super().__init__(name=name, parent=parent) + self.good_chan_parameter = ManualParameter( + "good_chan_parameter", instrument=self + ) + self.bad_chan_parameter = ManualParameter("bad_chan_parameter") + + +class DummyInstrument(Instrument): + def __init__(self, name: str): + super().__init__(name=name) + self.good_inst_parameter = ManualParameter( + "good_inst_parameter", instrument=self + ) + self.bad_inst_parameter = ManualParameter("bad_inst_parameter") + self.module = DummyModule(name="module", parent=self) + + +class DummyDelegateInstrument(InstrumentBase): + def __init__(self, name: str): + super().__init__(name=name) + self.inst_delegate = DelegateParameter( + name="inst_delegate", source=None, instrument=self, bind_to_instrument=True + ) + self.module = DummyDelegateModule(name="dummy_delegate_module", parent=self) + self.inst_base_parameter = ManualParameter( + "inst_base_parameter", instrument=self + ) + + +class DummyDelegateModule(InstrumentModule): + def __init__(self, name: str, parent: InstrumentBase): + super().__init__(name=name, parent=parent) + self.chan_delegate = DelegateParameter( + name="chan_delegate", source=None, instrument=self, bind_to_instrument=True + ) + + +class UserLinkingParameter(Parameter): + def __init__( + self, name: str, linked_parameter: Parameter | None = None, **kwargs: Any + ): + super().__init__(name=name, **kwargs) + self.linked_parameter: Parameter | None = linked_parameter + + +@pytest.fixture(name="instrument_fixture") +def make_instrument_fixture(): + inst = DummyInstrument("dummy_instrument") + InferAttrs.clear() + try: + yield inst + finally: + inst.close() + + +@pytest.fixture(name="good_inst_delegates") +def make_good_delegate_parameters(instrument_fixture): + inst = instrument_fixture + 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) + 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 + + assert get_root_parameter(good_inst_del_1) is inst.good_inst_parameter + assert get_root_parameter(good_inst_del_2) is inst.good_inst_parameter + + assert ( + get_root_parameter(good_inst_del_3, "linked_parameter") + is inst.good_inst_parameter + ) + + InferAttrs.clear() + assert get_root_parameter(good_inst_del_3) is good_inst_del_3 + + InferAttrs.add("linked_parameter") + assert get_root_parameter(good_inst_del_3) is inst.good_inst_parameter + + +def test_get_root_parameter_no_source(good_inst_delegates): + good_inst_del_1, good_inst_del_2, _ = good_inst_delegates + + good_inst_del_1.source = None + + with pytest.raises(InferError) as exc_info: + get_root_parameter(good_inst_del_2) + assert "is not attached to a source" in str(exc_info.value) + + +def test_get_root_parameter_no_user_attr(good_inst_delegates): + _, _, good_inst_del_3 = good_inst_delegates + InferAttrs.clear() + assert get_root_parameter(good_inst_del_3, "external_parameter") is good_inst_del_3 + + +def test_get_root_parameter_none_user_attr(good_inst_delegates): + _, _, good_inst_del_3 = good_inst_delegates + good_inst_del_3.linked_parameter = None + with pytest.raises(InferError) as exc_info: + get_root_parameter(good_inst_del_3, "linked_parameter") + assert "is not attached to a source on attribute" in str(exc_info.value) + + +def test_infer_instrument_valid(instrument_fixture, good_inst_delegates): + inst = instrument_fixture + _, _, good_inst_del_3 = good_inst_delegates + InferAttrs.add("linked_parameter") + assert infer_instrument(good_inst_del_3) is inst + + +def test_infer_instrument_no_instrument(instrument_fixture): + inst = instrument_fixture + no_inst_delegate = DelegateParameter( + "no_inst_delegate", source=inst.bad_inst_parameter + ) + with pytest.raises(InferError) as exc_info: + infer_instrument(no_inst_delegate) + assert "has no instrument" in str(exc_info.value) + + +def test_infer_instrument_root_instrument_base(): + delegate_inst = DummyDelegateInstrument("dummy_delegate_instrument") + + with pytest.raises(InferError) as exc_info: + infer_instrument(delegate_inst.inst_base_parameter) + assert "Could not determine source instrument for parameter" in str(exc_info.value) + + +def test_infer_channel_valid(instrument_fixture): + inst = instrument_fixture + chan_delegate = DelegateParameter( + "chan_delegate", source=inst.module.good_chan_parameter + ) + assert infer_channel(chan_delegate) is inst.module + + +def test_infer_channel_no_channel(instrument_fixture): + inst = instrument_fixture + no_chan_delegate = DelegateParameter( + "no_chan_delegate", source=inst.module.bad_chan_parameter + ) + with pytest.raises(InferError) as exc_info: + infer_channel(no_chan_delegate) + assert "has no instrument" in str(exc_info.value) + + inst_but_not_chan_delegate = DelegateParameter( + "inst_but_not_chan_delegate", source=inst.good_inst_parameter + ) + with pytest.raises(InferError) as exc_info: + infer_channel(inst_but_not_chan_delegate) + assert "Could not determine a root instrument channel" in str(exc_info.value) + + +def test_get_parameter_chain(instrument_fixture, good_inst_delegates): + inst = instrument_fixture + good_inst_del_1, good_inst_del_2, good_inst_del_3 = good_inst_delegates + parameter_chain = get_parameter_chain(good_inst_del_3, "linked_parameter") + expected_chain = ( + good_inst_del_3, + good_inst_del_2, + good_inst_del_1, + inst.good_inst_parameter, + ) + assert np.all( + [parameter_chain[i] is param for i, param in enumerate(expected_chain)] + ) + + # This is a broken chain. get_root_parameter would throw an InferError, but + # get_parameter_chain should run successfully + good_inst_del_1.source = None + parameter_chain = get_parameter_chain(good_inst_del_3, "linked_parameter") + expected_chain = ( + good_inst_del_3, + good_inst_del_2, + good_inst_del_1, + ) + assert np.all( + [parameter_chain[i] is param for i, param in enumerate(expected_chain)] + ) + + # Make the linked_parameter at the end of the chain + good_inst_del_3.linked_parameter = None + good_inst_del_1.source = good_inst_del_3 + parameter_chain = get_parameter_chain(good_inst_del_2, "linked_parameter") + expected_chain = ( + good_inst_del_2, + good_inst_del_1, + good_inst_del_3, + ) + assert np.all( + [parameter_chain[i] is param for i, param in enumerate(expected_chain)] + ) + + +def test_parameters_on_delegate_instruments(instrument_fixture, good_inst_delegates): + inst = instrument_fixture + _, good_inst_del_2, _ = good_inst_delegates + + delegate_inst = DummyDelegateInstrument("dummy_delegate_instrument") + delegate_inst.inst_delegate.source = good_inst_del_2 + delegate_inst.module.chan_delegate.source = inst.module.good_chan_parameter + + assert infer_channel(delegate_inst.module.chan_delegate) is inst.module + assert infer_instrument(delegate_inst.module.chan_delegate) is inst + assert infer_instrument(delegate_inst.inst_delegate) is inst + + +def test_merge_user_and_class_attrs(): + InferAttrs.add("attr1") + attr_set = _merge_user_and_class_attrs("attr2") + assert set(("attr1", "attr2")) == attr_set + + attr_set_list = _merge_user_and_class_attrs(("attr2", "attr3")) + assert set(("attr1", "attr2", "attr3")) == attr_set_list + + +def test_infer_attrs(): + InferAttrs.clear() + assert InferAttrs.known_attrs() == () + + InferAttrs.add("attr1") + assert set(InferAttrs.known_attrs()) == set(("attr1",)) + + InferAttrs.add("attr2") + InferAttrs.discard("attr1") + assert set(InferAttrs.known_attrs()) == set(("attr2",)) + + InferAttrs.add(("attr1", "attr3")) + assert set(InferAttrs.known_attrs()) == set(("attr1", "attr2", "attr3")) + + +def test_get_parameter_chain_with_loops(good_inst_delegates): + good_inst_del_1, good_inst_del_2, good_inst_del_3 = good_inst_delegates + good_inst_del_1.source = good_inst_del_3 + parameter_chain = get_parameter_chain(good_inst_del_3, "linked_parameter") + expected_chain = ( + good_inst_del_3, + good_inst_del_2, + good_inst_del_1, + good_inst_del_3, + ) + assert np.all( + [parameter_chain[i] is param for i, param in enumerate(expected_chain)] + ) + + +def test_get_root_parameter_with_loops(good_inst_delegates): + good_inst_del_1, good_inst_del_2, good_inst_del_3 = good_inst_delegates + good_inst_del_1.source = good_inst_del_3 + 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)