Merge pull request #7905 from jenshnielsen/jenshnielsen/strict_keysight

Keysight drivers avoid internal dependency on dynamic attributes

Author: jenshnielsen

src/qcodes/instrument_drivers/Keysight/Infiniium.py

@@ -13,6 +13,7 @@
 import qcodes.validators as vals
 from qcodes.instrument import (
     ChannelList,
+    ChannelTuple,
     InstrumentBase,
     InstrumentBaseKWArgs,
     InstrumentChannel,
@@ -89,7 +90,11 @@ def get_raw(self) -> npt.NDArray:
         )
 
 
-class DSOTraceParam(ParameterWithSetpoints):
+class DSOTraceParam(
+    ParameterWithSetpoints[
+        npt.NDArray, "KeysightInfiniiumChannel | KeysightInfiniiumFunction"
+    ]
+):
     """
     Trace parameter for the Infiniium series DSO
     """
@@ -180,7 +185,7 @@ def update_setpoints(self, preamble: "Sequence[str] | None" = None) -> None:
         acquisition if instr.cache_setpoints is False
         """
         instrument: KeysightInfiniiumChannel | KeysightInfiniiumFunction
-        instrument = self.instrument  # type: ignore[assignment]
+        instrument = self.instrument
         if preamble is None:
             instrument.write(f":WAV:SOUR {self._channel}")
             preamble = instrument.ask(":WAV:PRE?").strip().split(",")
@@ -257,13 +262,15 @@ def __init__(
         self,
         parent: InstrumentBase,
         name: str,
+        channel: str,
         **kwargs: "Unpack[InstrumentBaseKWArgs]",
     ) -> None:
         """
         Add parameters to measurement subsystem. Note: This should not be initialized
         directly, rather initialize BoundMeasurementSubsystem
         or UnboundMeasurementSubsystem.
         """
+        self._channel = channel
         super().__init__(parent, name, **kwargs)
 
         ###################################
@@ -477,11 +484,8 @@ def __init__(
         """
         Initialize measurement subsystem bound to a specific channel
         """
-        # Bind the channel
-        self._channel = parent.channel_name
-
         # Initialize measurement parameters
-        super().__init__(parent, name, **kwargs)
+        super().__init__(parent, name, channel=parent.channel_name, **kwargs)
 
 
 BoundMeasurement = KeysightInfiniiumBoundMeasurement
@@ -500,11 +504,8 @@ def __init__(
         """
         Initialize measurement subsystem where target is set by the parameter `source`.
         """
-        # Blank channel
-        self._channel = ""
-
         # Initialize measurement parameters
-        super().__init__(parent, name, **kwargs)
+        super().__init__(parent, name, channel="", **kwargs)
 
         self.source = Parameter(
             name="source",
@@ -515,6 +516,12 @@ def __init__(
             snapshot_value=False,
         )
 
+    @property
+    def root_instrument(self) -> "KeysightInfiniium":
+        root_instrument = super().root_instrument
+        assert isinstance(root_instrument, KeysightInfiniium)
+        return root_instrument
+
     def _validate_source(self, source: str) -> str:
         """Validate and set the source."""
         valid_channels = f"CHAN[1-{self.root_instrument.no_channels}]"
@@ -692,7 +699,7 @@ def _get_func(self) -> str:
 """
 
 
-class KeysightInfiniiumChannel(InstrumentChannel):
+class KeysightInfiniiumChannel(InstrumentChannel["KeysightInfiniium"]):
     def __init__(
         self,
         parent: "KeysightInfiniium",
@@ -1081,7 +1088,10 @@ def __init__(
             channel = KeysightInfiniiumChannel(self, f"chan{i}", i)
             _channels.append(channel)
             self.add_submodule(f"ch{i}", channel)
-        self.add_submodule("channels", _channels.to_channel_tuple())
+        self.channels: ChannelTuple[KeysightInfiniiumChannel] = self.add_submodule(
+            "channels", _channels.to_channel_tuple()
+        )
+        """Tuple of oscilloscope channels."""
 
         # Functions
         _functions = ChannelList(
@@ -1093,11 +1103,17 @@ def __init__(
             self.add_submodule(f"func{i}", function)
         # Have to call channel list "funcs" here as functions is a
         # reserved name in Instrument.
-        self.add_submodule("funcs", _functions.to_channel_tuple())
+        self.funcs: ChannelTuple[KeysightInfiniiumFunction] = self.add_submodule(
+            "funcs", _functions.to_channel_tuple()
+        )
+        """Tuple of oscilloscope functions."""
 
         # Submodules
         meassubsys = KeysightInfiniiumUnboundMeasurement(self, "measure")
-        self.add_submodule("measure", meassubsys)
+        self.measure: KeysightInfiniiumUnboundMeasurement = self.add_submodule(
+            "measure", meassubsys
+        )
+        """Unbound measurement subsystem."""
 
     def _query_capabilities(self) -> None:
         """

src/qcodes/instrument_drivers/Keysight/KeysightAgilent_33XXX.py

@@ -4,7 +4,6 @@
 
 from qcodes import validators as vals
 from qcodes.instrument import (
-    Instrument,
     InstrumentBaseKWArgs,
     InstrumentChannel,
     VisaInstrument,
@@ -26,14 +25,14 @@
 # 33200, 33500, and 33600
 
 
-class Keysight33xxxOutputChannel(InstrumentChannel):
+class Keysight33xxxOutputChannel(InstrumentChannel["Keysight33xxx"]):
     """
     Class to hold the output channel of a Keysight 33xxxx waveform generator.
     """
 
     def __init__(
         self,
-        parent: Instrument,
+        parent: "Keysight33xxx",
         name: str,
         channum: int,
         **kwargs: "Unpack[InstrumentBaseKWArgs]",
@@ -317,15 +316,15 @@ def val_parser(parser: type, inputstring: str) -> float | int:
 OutputChannel = Keysight33xxxOutputChannel
 
 
-class Keysight33xxxSyncChannel(InstrumentChannel):
+class Keysight33xxxSyncChannel(InstrumentChannel["Keysight33xxx"]):
     """
     Class to hold the sync output of a Keysight 33xxxx waveform generator.
     Has very few parameters for single channel instruments.
     """
 
     def __init__(
         self,
-        parent: Instrument,
+        parent: "Keysight33xxx",
         name: str,
         **kwargs: "Unpack[InstrumentBaseKWArgs]",
     ):

src/qcodes/instrument_drivers/Keysight/Keysight_N9030B.py

@@ -98,7 +98,6 @@ def __init__(
         **kwargs: Unpack[InstrumentBaseKWArgs],
     ):
         super().__init__(parent, name, *arg, **kwargs)
-        self.root_instrument: KeysightN9030B
 
         self._additional_wait = additional_wait
         self._min_freq = -8e7
@@ -111,7 +110,7 @@ def __init__(
         }
         opt: str | None = None
         for hw_opt_for_max_freq in self._valid_max_freq:
-            if hw_opt_for_max_freq in self.root_instrument.options():
+            if hw_opt_for_max_freq in self.parent.options():
                 opt = hw_opt_for_max_freq
         assert opt is not None
         self._max_freq = self._valid_max_freq[opt]
@@ -454,23 +453,22 @@ def _get_data(self, trace_num: int) -> npt.NDArray[np.float64]:
         """
         Gets data from the measurement.
         """
-        root_instr = self.root_instrument
         # Check if we should run a new sweep
-        auto_sweep = root_instr.auto_sweep()
+        auto_sweep = self.parent.auto_sweep()
 
         if auto_sweep:
             # If we need to run a sweep, we need to set the timeout to take into account
             # the sweep time
             timeout = self.sweep_time() + self._additional_wait
-            with root_instr.timeout.set_to(timeout):
-                data = root_instr.visa_handle.query_binary_values(
-                    f":READ:{root_instr.measurement()}{trace_num}?",
+            with self.parent.timeout.set_to(timeout):
+                data = self.parent.visa_handle.query_binary_values(
+                    f":READ:{self.parent.measurement()}{trace_num}?",
                     datatype="d",
                     is_big_endian=False,
                 )
         else:
-            data = root_instr.visa_handle.query_binary_values(
-                f":FETC:{root_instr.measurement()}{trace_num}?",
+            data = self.parent.visa_handle.query_binary_values(
+                f":FETC:{self.parent.measurement()}{trace_num}?",
                 datatype="d",
                 is_big_endian=False,
             )
@@ -491,9 +489,9 @@ def setup_swept_sa_sweep(self, start: float, stop: float, npts: int) -> None:
         """
         Sets up the Swept SA measurement sweep for Spectrum Analyzer Mode.
         """
-        self.root_instrument.mode("SA")
-        if "SAN" in self.root_instrument.available_meas():
-            self.root_instrument.measurement("SAN")
+        self.parent.mode("SA")
+        if "SAN" in self.parent.available_meas():
+            self.parent.measurement("SAN")
         else:
             raise RuntimeError(
                 "Swept SA measurement is not available on your "
@@ -537,7 +535,7 @@ def __init__(
         }
         opt: str | None = None
         for hw_opt_for_max_freq in self._valid_max_freq:
-            if hw_opt_for_max_freq in self.root_instrument.options():
+            if hw_opt_for_max_freq in self.parent.options():
                 opt = hw_opt_for_max_freq
         assert opt is not None
         self._max_freq = self._valid_max_freq[opt]
@@ -668,9 +666,8 @@ def _get_data(self, trace_num: int) -> ParamRawDataType:
         """
         Gets data from the measurement.
         """
-        root_instr = self.root_instrument
-        measurement = root_instr.measurement()
-        raw_data = root_instr.visa_handle.query_binary_values(
+        measurement = self.parent.measurement()
+        raw_data = self.parent.visa_handle.query_binary_values(
             f":READ:{measurement}1?",
             datatype="d",
             is_big_endian=False,
@@ -684,26 +681,26 @@ def _get_data(self, trace_num: int) -> ParamRawDataType:
             return -1 * np.ones(self.npts())
 
         try:
-            data = root_instr.visa_handle.query_binary_values(
+            data = self.parent.visa_handle.query_binary_values(
                 f":READ:{measurement}{trace_num}?",
                 datatype="d",
                 is_big_endian=False,
             )
-            data = np.array(data).reshape((-1, 2))
+            data_array = np.array(data).reshape((-1, 2))
         except TimeoutError as e:
             raise TimeoutError("Couldn't receive any data. Command timed out.") from e
 
-        return data[:, 1]
+        return data_array[:, 1]
 
     def setup_log_plot_sweep(
         self, start_offset: float, stop_offset: float, npts: int
     ) -> None:
         """
         Sets up the Log Plot measurement sweep for Phase Noise Mode.
         """
-        self.root_instrument.mode("PNOISE")
-        if "LPL" in self.root_instrument.available_meas():
-            self.root_instrument.measurement("LPL")
+        self.parent.mode("PNOISE")
+        if "LPL" in self.parent.available_meas():
+            self.parent.measurement("LPL")
         else:
             raise RuntimeError(
                 "Log Plot measurement is not available on your "

src/qcodes/instrument_drivers/Keysight/N52xx.py

@@ -8,6 +8,7 @@
 
 from qcodes.instrument import (
     ChannelList,
+    ChannelTuple,
     InstrumentBaseKWArgs,
     InstrumentChannel,
     VisaInstrument,
@@ -69,7 +70,7 @@ def get_raw(self) -> npt.NDArray:
         return np.linspace(0, self._stopparam(), self._pointsparam())
 
 
-class FormattedSweep(ParameterWithSetpoints):
+class FormattedSweep(ParameterWithSetpoints[npt.NDArray, "KeysightPNATrace"]):
     """
     Mag will run a sweep, including averaging, before returning data.
     As such, wait time in a loop is not needed.
@@ -96,7 +97,7 @@ def setpoints(self) -> "Sequence[ParameterBase]":
         """
         if self.instrument is None:
             raise RuntimeError("Cannot return setpoints if not attached to instrument")
-        root_instrument: KeysightPNABase = self.root_instrument  # type: ignore[assignment]
+        root_instrument: KeysightPNABase = self.root_instrument
         sweep_type = root_instrument.sweep_type()
         if sweep_type == "LIN":
             return (root_instrument.frequency_axis,)
@@ -115,10 +116,16 @@ def setpoints(self, setpoints: Any) -> None:
         """
         return
 
+    @property
+    def root_instrument(self) -> "KeysightPNABase":
+        root_instrument = super().root_instrument
+        assert isinstance(root_instrument, KeysightPNABase)
+        return root_instrument
+
     def get_raw(self) -> npt.NDArray:
         if self.instrument is None:
             raise RuntimeError("Cannot get data without instrument")
-        root_instr = self.instrument.root_instrument
+        root_instr = self.root_instrument
         # Check if we should run a new sweep
         auto_sweep = root_instr.auto_sweep()
 
@@ -130,12 +137,12 @@ def get_raw(self) -> npt.NDArray:
         data = root_instr.visa_handle.query_binary_values(
             "CALC:DATA? FDATA", datatype="f", is_big_endian=True
         )
-        data = np.array(data)
+        data_array = np.array(data)
         # Restore previous state if it was changed
         if auto_sweep:
             root_instr.sweep_mode(prev_mode)
 
-        return data
+        return data_array
 
 
 class KeysightPNAPort(InstrumentChannel):
@@ -182,7 +189,7 @@ def _set_power_limits(self, min_power: float, max_power: float) -> None:
 "Alis for backwards compatibility"
 
 
-class KeysightPNATrace(InstrumentChannel):
+class KeysightPNATrace(InstrumentChannel["KeysightPNABase"]):
     """
     Allow operations on individual PNA traces.
     """
@@ -292,6 +299,12 @@ def __init__(
         )
         """Parameter polar"""
 
+    @property
+    def root_instrument(self) -> "KeysightPNABase":
+        root_instrument = super().root_instrument
+        assert isinstance(root_instrument, KeysightPNABase)
+        return root_instrument
+
     def disable(self) -> None:
         """
         Disable this trace on the PNA
@@ -438,7 +451,11 @@ def __init__(
             )
             ports.append(port)
             self.add_submodule(f"port{port_num}", port)
-        self.add_submodule("ports", ports.to_channel_tuple())
+
+        self.ports: ChannelTuple[KeysightPNAPort] = self.add_submodule(
+            "ports", ports.to_channel_tuple()
+        )
+        """Tuple of KeysightPNAPort submodules"""
 
         # RF output
         self.output: Parameter = self.add_parameter(

src/qcodes/instrument_drivers/Keysight/keysight_34934a.py

@@ -13,11 +13,11 @@
 
     from qcodes.instrument import (
         InstrumentBaseKWArgs,
-        InstrumentChannel,
-        VisaInstrument,
     )
     from qcodes.parameters import Parameter
 
+    from .keysight_34980a import Keysight34980A
+
 
 class Keysight34934A(Keysight34980ASwitchMatrixSubModule):
     """
@@ -32,7 +32,7 @@ class Keysight34934A(Keysight34980ASwitchMatrixSubModule):
 
     def __init__(
         self,
-        parent: "VisaInstrument | InstrumentChannel",
+        parent: "Keysight34980A",
         name: str,
         slot: int,
         **kwargs: "Unpack[InstrumentBaseKWArgs]",