[wip] Generic instrument in parameter

src/qcodes/instrument_drivers/AlazarTech/utils.py

@@ -5,15 +5,16 @@
 :mod:`.AlazarTech.helpers` module).
 """
 
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any
 
 from qcodes.parameters import Parameter, ParamRawDataType
 
 if TYPE_CHECKING:
-    from .ATS import AlazarTechATS
+    # ruff does not detect that this is used as a generic parameter below
+    from .ATS import AlazarTechATS  # noqa: F401
 
 
-class TraceParameter(Parameter):
+class TraceParameter(Parameter[Any, "AlazarTechATS"]):
     """
     A parameter that keeps track of if its value has been synced to
     the ``Instrument``. To achieve that, this parameter sets
@@ -38,6 +39,5 @@ def synced_to_card(self) -> bool:
         return self._synced_to_card
 
     def set_raw(self, value: ParamRawDataType) -> None:
-        instrument = cast("AlazarTechATS", self.instrument)
-        instrument._parameters_synced = False
+        self.instrument._parameters_synced = False
         self._synced_to_card = False

src/qcodes/instrument_drivers/Keysight/Keysight_N9030B.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar
 
 import numpy as np
+import numpy.typing as npt
 
 from qcodes.instrument import (
     InstrumentBaseKWArgs,
@@ -23,56 +24,62 @@
 
     from typing_extensions import Unpack
 
+_T = TypeVar(
+    "_T", bound="KeysightN9030BSpectrumAnalyzerMode | KeysightN9030BPhaseNoiseMode"
+)
+_S = TypeVar("_S")
+
 
-class FrequencyAxis(Parameter):
+class FrequencyAxis(
+    Parameter[
+        npt.NDArray[np.float64],
+        _T,
+    ],
+    Generic[_T],
+):
     def __init__(
         self,
-        start: Parameter,
-        stop: Parameter,
-        npts: Parameter,
+        start: Parameter[float, _T],
+        stop: Parameter[float, _T],
+        npts: Parameter[int, _T],
         *args: Any,
         **kwargs: Any,
     ) -> None:
         super().__init__(*args, **kwargs)
-        self._start: Parameter = start
-        self._stop: Parameter = stop
-        self._npts: Parameter = npts
+        self._start = start
+        self._stop = stop
+        self._npts = npts
 
     def get_raw(self) -> ParamRawDataType:
         start_val = self._start()
         stop_val = self._stop()
         npts_val = self._npts()
-        assert start_val is not None
-        assert stop_val is not None
-        assert npts_val is not None
+        if start_val is None or stop_val is None or npts_val is None:
+            raise RuntimeError("Start, Stop and Npts parameters must be set.")
         return np.linspace(start_val, stop_val, npts_val)
 
 
-class Trace(ParameterWithSetpoints):
+class Trace(ParameterWithSetpoints[_S, _T], Generic[_S, _T]):
     def __init__(
         self,
         number: int,
         *args: Any,
-        get_data: Callable[[int], ParamRawDataType],
+        get_data: Callable[[int], _S],
         **kwargs: Any,
     ) -> None:
         super().__init__(*args, **kwargs)
-        # the parameter classes should ideally be generic in instrument
-        # and root instrument classes so we can specialize here.
-        # for now we have to ignore a type error from pyright
-        self.instrument: (
-            KeysightN9030BSpectrumAnalyzerMode | KeysightN9030BPhaseNoiseMode
-        )
+        # while the parameter classes should ideally be generic in instrument
+        # type it is not generic in the root instrument type
         self.root_instrument: KeysightN9030B
 
         self.number = number
         self.get_data = get_data
 
-    def get_raw(self) -> ParamRawDataType:
+    def get_raw(self) -> _S:
         return self.get_data(self.number)
 
 
-class KeysightN9030BSpectrumAnalyzerMode(InstrumentChannel):
+class KeysightN9030BSpectrumAnalyzerMode(InstrumentChannel["KeysightN9030B"]):
     """
     Spectrum Analyzer Mode for Keysight N9030B instrument.
     """
@@ -86,6 +93,7 @@ def __init__(
         **kwargs: Unpack[InstrumentBaseKWArgs],
     ):
         super().__init__(parent, name, *arg, **kwargs)
+        self.root_instrument: KeysightN9030B
 
         self._additional_wait = additional_wait
         self._min_freq = -8e7
@@ -104,68 +112,82 @@ def __init__(
         self._max_freq = self._valid_max_freq[opt]
 
         # Frequency Parameters
-        self.start: Parameter = self.add_parameter(
-            name="start",
-            unit="Hz",
-            get_cmd=":SENSe:FREQuency:STARt?",
-            set_cmd=self._set_start,
-            get_parser=float,
-            vals=Numbers(self._min_freq, self._max_freq - 10),
-            docstring="Start Frequency",
+        self.start: Parameter[float, KeysightN9030BSpectrumAnalyzerMode] = (
+            self.add_parameter(
+                name="start",
+                unit="Hz",
+                get_cmd=":SENSe:FREQuency:STARt?",
+                set_cmd=self._set_start,
+                get_parser=float,
+                vals=Numbers(self._min_freq, self._max_freq - 10),
+                docstring="Start Frequency",
+            )
         )
         """Start Frequency"""
-        self.stop: Parameter = self.add_parameter(
-            name="stop",
-            unit="Hz",
-            get_cmd=":SENSe:FREQuency:STOP?",
-            set_cmd=self._set_stop,
-            get_parser=float,
-            vals=Numbers(self._min_freq + 10, self._max_freq),
-            docstring="Stop Frequency",
+        self.stop: Parameter[float, KeysightN9030BSpectrumAnalyzerMode] = (
+            self.add_parameter(
+                name="stop",
+                unit="Hz",
+                get_cmd=":SENSe:FREQuency:STOP?",
+                set_cmd=self._set_stop,
+                get_parser=float,
+                vals=Numbers(self._min_freq + 10, self._max_freq),
+                docstring="Stop Frequency",
+            )
         )
         """Stop Frequency"""
-        self.center: Parameter = self.add_parameter(
-            name="center",
-            unit="Hz",
-            get_cmd=":SENSe:FREQuency:CENTer?",
-            set_cmd=self._set_center,
-            get_parser=float,
-            vals=Numbers(self._min_freq + 5, self._max_freq - 5),
-            docstring="Sets and gets center frequency",
+        self.center: Parameter[float, KeysightN9030BSpectrumAnalyzerMode] = (
+            self.add_parameter(
+                name="center",
+                unit="Hz",
+                get_cmd=":SENSe:FREQuency:CENTer?",
+                set_cmd=self._set_center,
+                get_parser=float,
+                vals=Numbers(self._min_freq + 5, self._max_freq - 5),
+                docstring="Sets and gets center frequency",
+            )
         )
         """Sets and gets center frequency"""
-        self.span: Parameter = self.add_parameter(
-            name="span",
-            unit="Hz",
-            get_cmd=":SENSe:FREQuency:SPAN?",
-            set_cmd=self._set_span,
-            get_parser=float,
-            vals=Numbers(10, self._max_freq - self._min_freq),
-            docstring="Changes span of frequency",
+        self.span: Parameter[float, KeysightN9030BSpectrumAnalyzerMode] = (
+            self.add_parameter(
+                name="span",
+                unit="Hz",
+                get_cmd=":SENSe:FREQuency:SPAN?",
+                set_cmd=self._set_span,
+                get_parser=float,
+                vals=Numbers(10, self._max_freq - self._min_freq),
+                docstring="Changes span of frequency",
+            )
         )
         """Changes span of frequency"""
-        self.npts: Parameter = self.add_parameter(
-            name="npts",
-            get_cmd=":SENSe:SWEep:POINts?",
-            set_cmd=":SENSe:SWEep:POINts {}",
-            get_parser=int,
-            vals=Ints(1, 20001),
-            docstring="Number of points for the sweep",
+        self.npts: Parameter[int, KeysightN9030BSpectrumAnalyzerMode] = (
+            self.add_parameter(
+                name="npts",
+                get_cmd=":SENSe:SWEep:POINts?",
+                set_cmd=":SENSe:SWEep:POINts {}",
+                get_parser=int,
+                vals=Ints(1, 20001),
+                docstring="Number of points for the sweep",
+            )
         )
         """Number of points for the sweep"""
 
         # Amplitude/Input Parameters
-        self.mech_attenuation: Parameter = self.add_parameter(
-            name="mech_attenuation",
-            unit="dB",
-            get_cmd=":SENS:POW:ATT?",
-            set_cmd=":SENS:POW:ATT {}",
-            get_parser=int,
-            vals=Ints(0, 70),
-            docstring="Internal mechanical attenuation",
+        self.mech_attenuation: Parameter[int, KeysightN9030BSpectrumAnalyzerMode] = (
+            self.add_parameter(
+                name="mech_attenuation",
+                unit="dB",
+                get_cmd=":SENS:POW:ATT?",
+                set_cmd=":SENS:POW:ATT {}",
+                get_parser=int,
+                vals=Ints(0, 70),
+                docstring="Internal mechanical attenuation",
+            )
         )
         """Internal mechanical attenuation"""
-        self.preamp: Parameter = self.add_parameter(
+        self.preamp: Parameter[
+            Literal["LOW", "FULL"], KeysightN9030BSpectrumAnalyzerMode
+        ] = self.add_parameter(
             name="preamp",
             get_cmd=":SENS:POW:GAIN:BAND?",
             set_cmd=":SENS:POW:GAIN:BAND {}",
@@ -354,20 +376,25 @@ def __init__(
         """Sets up sweep type. Possible options are 'fft' and 'sweep'."""
 
         # Array (Data) Parameters
-        self.freq_axis: FrequencyAxis = self.add_parameter(
-            name="freq_axis",
-            label="Frequency",
-            unit="Hz",
-            start=self.start,
-            stop=self.stop,
-            npts=self.npts,
-            vals=Arrays(shape=(self.npts.get_latest,)),
-            parameter_class=FrequencyAxis,
-            docstring="Creates frequency axis for the sweep from start, "
-            "stop and npts values.",
+        self.freq_axis: FrequencyAxis[KeysightN9030BSpectrumAnalyzerMode] = (
+            self.add_parameter(
+                name="freq_axis",
+                label="Frequency",
+                unit="Hz",
+                start=self.start,
+                stop=self.stop,
+                npts=self.npts,
+                vals=Arrays(shape=(self.npts.get_latest,)),
+                parameter_class=FrequencyAxis,
+                docstring="Creates frequency axis for the sweep from start, "
+                "stop and npts values.",
+            )
         )
+
         """Creates frequency axis for the sweep from start, stop and npts values."""
-        self.trace: Trace = self.add_parameter(
+        self.trace: Trace[
+            npt.NDArray[np.float64], KeysightN9030BSpectrumAnalyzerMode
+        ] = self.add_parameter(
             name="trace",
             label="Trace",
             unit="dB",
@@ -418,7 +445,7 @@ def _set_span(self, val: float) -> None:
         self.write(f":SENSe:FREQuency:SPAN {val}")
         self.update_trace()
 
-    def _get_data(self, trace_num: int) -> ParamRawDataType:
+    def _get_data(self, trace_num: int) -> npt.NDArray[np.float64]:
         """
         Gets data from the measurement.
         """
@@ -443,8 +470,8 @@ def _get_data(self, trace_num: int) -> ParamRawDataType:
                 is_big_endian=False,
             )
 
-        data = np.array(data).reshape((-1, 2))
-        return data[:, 1]
+        data_array = np.array(data).reshape((-1, 2))
+        return data_array[:, 1]
 
     def update_trace(self) -> None:
         """
@@ -481,7 +508,7 @@ def autotune(self) -> None:
         self.center()
 
 
-class KeysightN9030BPhaseNoiseMode(InstrumentChannel):
+class KeysightN9030BPhaseNoiseMode(InstrumentChannel["KeysightN9030B"]):
     """
     Phase Noise Mode for Keysight N9030B instrument.
     """

src/qcodes/instrument_drivers/Keysight/keysightb1500/KeysightB1517A.py

@@ -1,6 +1,6 @@
 import re
 import textwrap
-from typing import TYPE_CHECKING, Any, Literal, NotRequired, cast, overload
+from typing import TYPE_CHECKING, Any, Literal, NotRequired, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -701,7 +701,7 @@ def _get_sweep_steps_parser(response: str) -> SweepSteps:
 """
 
 
-class _ParameterWithStatus(Parameter):
+class _ParameterWithStatus(Parameter[Any, "KeysightB1517A"]):
     def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
 
@@ -728,7 +728,7 @@ def snapshot_base(
 
 class _SpotMeasurementVoltageParameter(_ParameterWithStatus):
     def set_raw(self, value: ParamRawDataType) -> None:
-        smu = cast("KeysightB1517A", self.instrument)
+        smu = self.instrument
 
         if smu._source_config["output_range"] is None:
             smu._source_config["output_range"] = constants.VOutputRange.AUTO
@@ -752,7 +752,7 @@ def set_raw(self, value: ParamRawDataType) -> None:
         )
 
     def get_raw(self) -> ParamRawDataType:
-        smu = cast("KeysightB1517A", self.instrument)
+        smu = self.instrument
 
         msg = MessageBuilder().tv(
             chnum=smu.channels[0],
@@ -769,7 +769,7 @@ def get_raw(self) -> ParamRawDataType:
 
 class _SpotMeasurementCurrentParameter(_ParameterWithStatus):
     def set_raw(self, value: ParamRawDataType) -> None:
-        smu = cast("KeysightB1517A", self.instrument)
+        smu = self.instrument
 
         if smu._source_config["output_range"] is None:
             smu._source_config["output_range"] = constants.IOutputRange.AUTO
@@ -793,7 +793,7 @@ def set_raw(self, value: ParamRawDataType) -> None:
         )
 
     def get_raw(self) -> ParamRawDataType:
-        smu = cast("KeysightB1517A", self.instrument)
+        smu = self.instrument
 
         msg = MessageBuilder().ti(
             chnum=smu.channels[0],