Merge pull request #7659 from jenshnielsen/generic_parent

Make InstrumentModule/Channel take a generic instrument argument

Author: jenshnielsen

docs/changes/newsfragments/7659.improved

@@ -0,0 +1 @@
+The ``InstrumentModule`` and its alias ``InstrumentChannel`` now take an optional generic argument allowing you to specify the type of the parent instrument for type checking.

docs/changes/newsfragments/7659.improved_driver

@@ -0,0 +1,2 @@
+The ``KeysightE4980A`` driver now names the ``correction`` submodule correctly as ``correction`` reflecting the public attribute to access the module.
+This also means that in the snapshot ``correction`` is used as the module name rather than ``_correction``

src/qcodes/instrument/channel.py

@@ -5,7 +5,9 @@
 import sys
 import warnings
 from collections.abc import Callable, Iterable, Iterator, MutableSequence, Sequence
-from typing import TYPE_CHECKING, Any, TypeVar, cast, overload
+from typing import TYPE_CHECKING, Any, Generic, cast, overload
+
+from typing_extensions import TypeVar
 
 from qcodes.metadatable import MetadatableWithName
 from qcodes.parameters import (
@@ -28,7 +30,12 @@
     from .instrument_base import InstrumentBaseKWArgs
 
 
-class InstrumentModule(InstrumentBase):
+_TIB_co = TypeVar(
+    "_TIB_co", bound="InstrumentBase", default=InstrumentBase, covariant=True
+)
+
+
+class InstrumentModule(InstrumentBase, Generic[_TIB_co]):
     """
     Base class for a module in an instrument.
     This could be in the form of a channel (e.g. something that
@@ -45,7 +52,7 @@ class InstrumentModule(InstrumentBase):
     """
 
     def __init__(
-        self, parent: InstrumentBase, name: str, **kwargs: Unpack[InstrumentBaseKWArgs]
+        self, parent: _TIB_co, name: str, **kwargs: Unpack[InstrumentBaseKWArgs]
     ) -> None:
         # need to specify parent before `super().__init__` so that the right
         # `full_name` is available in that scope. `full_name` is used for
@@ -76,7 +83,7 @@ def ask_raw(self, cmd: str) -> str:
         return self._parent.ask_raw(cmd)
 
     @property
-    def parent(self) -> InstrumentBase:
+    def parent(self) -> _TIB_co:
         return self._parent
 
     @property
@@ -90,7 +97,7 @@ def name_parts(self) -> list[str]:
         return name_parts
 
 
-class InstrumentChannel(InstrumentModule):
+class InstrumentChannel(InstrumentModule[_TIB_co], Generic[_TIB_co]):
     pass
 
 

src/qcodes/instrument_drivers/Keithley/Keithley_2450.py

@@ -213,7 +213,7 @@ class _FunctionMode(TypedDict):
     range_vals: Numbers
 
 
-class Keithley2450Sense(InstrumentChannel):
+class Keithley2450Sense(InstrumentChannel["Keithley2450"]):
     """
     The sense module of the Keithley 2450 SMU.
 
@@ -340,7 +340,7 @@ def _measure(self) -> float | str:
         return float(self.ask(f":MEASure? '{buffer_name}'"))
 
     def _measure_sweep(self) -> npt.NDArray:
-        source = cast("Keithley2450Source", self.parent.source)
+        source = self.parent.source
         source.sweep_start()
         buffer_name = self.parent.buffer_name()
         buffer = cast(

src/qcodes/instrument_drivers/Keysight/keysight_e4980a.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any
 
 from packaging import version
 from pyvisa.errors import VisaIOError
@@ -371,15 +371,14 @@ def __init__(
         )
         """This parameter tracks the signal mode which is being set."""
 
-        self.add_submodule("_correction", KeysightE4980ACorrection(self, "correction"))
+        self.correction: KeysightE4980ACorrection = self.add_submodule(
+            "correction", KeysightE4980ACorrection(self, "correction")
+        )
+        """Correction submodule"""
+
         self._set_signal_mode_on_driver_initialization()
         self.connect_message()
 
-    @property
-    def correction(self) -> KeysightE4980ACorrection:
-        submodule = self.submodules["_correction"]
-        return cast("KeysightE4980ACorrection", submodule)
-
     @property
     def measure_impedance(self) -> KeysightE4980AMeasurementPair:
         return self._get_complex_impedance()

src/qcodes/instrument_drivers/Lakeshore/Lakeshore_model_325.py

@@ -7,7 +7,6 @@
     SupportsBytes,
     SupportsIndex,
     TextIO,
-    cast,
 )
 
 from qcodes.instrument import (
@@ -303,7 +302,7 @@ def set_data(
             self.write(cmd_str)
 
 
-class LakeshoreModel325Sensor(InstrumentChannel):
+class LakeshoreModel325Sensor(InstrumentChannel["LakeshoreModel325"]):
     """
     InstrumentChannel for a single sensor of a Lakeshore Model 325.
 
@@ -396,8 +395,7 @@ def decode_sensor_status(sum_of_codes: int) -> str:
 
     @property
     def curve(self) -> LakeshoreModel325Curve:
-        parent = cast("LakeshoreModel325", self.parent)
-        return LakeshoreModel325Curve(parent, self.curve_index())
+        return LakeshoreModel325Curve(self.parent, self.curve_index())
 
 
 class LakeshoreModel325Heater(InstrumentChannel):
@@ -582,7 +580,10 @@ def __init__(
         super().__init__(name, address, **kwargs)
 
         sensors = ChannelList(
-            self, "sensor", LakeshoreModel325Sensor, snapshotable=False
+            self,
+            "sensor",
+            LakeshoreModel325Sensor,
+            snapshotable=False,
         )
 
         self.sensor_A: LakeshoreModel325Sensor = self.add_submodule(

src/qcodes/instrument_drivers/rigol/Rigol_DG1062.py

@@ -1,15 +1,15 @@
 import logging
 from functools import partial
-from typing import TYPE_CHECKING, Any, ClassVar, cast
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from qcodes import validators as vals
 from qcodes.instrument import (
-    ChannelList,
     InstrumentBaseKWArgs,
     InstrumentChannel,
     VisaInstrument,
     VisaInstrumentKWArgs,
 )
+from qcodes.instrument.channel import ChannelTuple
 from qcodes.utils import partial_with_docstring
 
 if TYPE_CHECKING:
@@ -20,7 +20,7 @@
 log = logging.getLogger(__name__)
 
 
-class RigolDG1062Burst(InstrumentChannel):
+class RigolDG1062Burst(InstrumentChannel["RigolDG1062"]):
     """
     Burst commands for the DG1062. We make a separate channel for these to
     group burst commands together.
@@ -129,7 +129,7 @@ def trigger(self) -> None:
         self.parent.write_raw(f":SOUR{self.channel}:BURS:TRIG")
 
 
-class RigolDG1062Channel(InstrumentChannel):
+class RigolDG1062Channel(InstrumentChannel["RigolDG1062"]):
     min_impedance = 1
     max_impedance = 10000
 
@@ -261,10 +261,9 @@ def __init__(
         For other waveforms it will give the user an error
         """
 
-        burst = RigolDG1062Burst(
-            cast("RigolDG1062", self.parent), "burst", self.channel
-        )
-        self.add_submodule("burst", burst)
+        burst = RigolDG1062Burst(self.parent, "burst", self.channel)
+        self.burst: RigolDG1062Burst = self.add_submodule("burst", burst)
+        """"Burst submodule"""
 
         # We want to be able to do the following:
         # >>> help(gd.channels[0].sin)
@@ -371,7 +370,7 @@ def _set_waveform_params(self, **params_dict: float) -> None:
         string += ",".join(values)
         self.parent.write_raw(string)
 
-    def _get_duty_cycle(self) -> float:
+    def _get_duty_cycle(self) -> str:
         """
         Reads the duty cycle after checking waveform
         """
@@ -415,13 +414,20 @@ def __init__(
     ):
         super().__init__(name, address, **kwargs)
 
-        channels = ChannelList(self, "channel", RigolDG1062Channel, snapshotable=False)
-
-        for ch_num in [1, 2]:
-            ch_name = f"ch{ch_num}"
-            channel = RigolDG1062Channel(self, ch_name, ch_num)
-            channels.append(channel)
-            self.add_submodule(ch_name, channel)
-
-        self.add_submodule("channels", channels.to_channel_tuple())
+        self.ch1 = self.add_submodule("ch1", RigolDG1062Channel(self, "ch1", 1))
+        """Channel 1 submodule"""
+        self.ch2 = self.add_submodule("ch2", RigolDG1062Channel(self, "ch2", 2))
+        """Channel 2 submodule"""
+
+        self.channels: ChannelTuple[RigolDG1062Channel] = self.add_submodule(
+            "channels",
+            ChannelTuple(
+                self,
+                "channel",
+                chan_type=RigolDG1062Channel,
+                chan_list=(self.ch1, self.ch2),
+                snapshotable=False,
+            ),
+        )
+        """Tuple of channels"""
         self.connect_message()