Make Parameter classes generic in the instrument type

Author: jenshnielsen

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["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/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["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],

src/qcodes/instrument_drivers/tektronix/DPO7200xx.py

@@ -7,7 +7,7 @@
 import textwrap
 import time
 from functools import partial
-from typing import TYPE_CHECKING, Any, ClassVar, cast
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import numpy as np
 import numpy.typing as npt
@@ -767,7 +767,7 @@ def _trigger_type(self, value: str) -> None:
         self.write(f"TRIGger:{self._identifier}:TYPE {value}")
 
 
-class TektronixDPOMeasurementParameter(Parameter):
+class TektronixDPOMeasurementParameter(Parameter["TektronixDPOMeasurement"]):
     """
     A measurement parameter does not only return the instantaneous value
     of a measurement, but can also return some statistics. The accumulation
@@ -778,7 +778,7 @@ class TektronixDPOMeasurementParameter(Parameter):
     """
 
     def _get(self, metric: str) -> float:
-        measurement_channel = cast("TektronixDPOMeasurement", self.instrument)
+        measurement_channel = self.instrument
         if measurement_channel.type.get_latest() != self.name:
             measurement_channel.type(self.name)
 

src/qcodes/parameters/parameter.py

@@ -6,7 +6,9 @@
 import logging
 import os
 from types import MethodType
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Generic, Literal
+
+from typing_extensions import TypeVar
 
 from .command import Command
 from .parameter_base import ParamDataType, ParameterBase, ParamRawDataType
@@ -24,7 +26,15 @@
 log = logging.getLogger(__name__)
 
 
-class Parameter(ParameterBase):
+_InstrumentType_co = TypeVar(
+    "_InstrumentType_co",
+    bound="InstrumentBase | None",
+    default="InstrumentBase | None",
+    covariant=True,
+)
+
+
+class Parameter(ParameterBase[_InstrumentType_co], Generic[_InstrumentType_co]):
     """
     A parameter represents a single degree of freedom. Most often,
     this is the standard parameter for Instruments, though it can also be
@@ -172,7 +182,7 @@ class Parameter(ParameterBase):
     def __init__(
         self,
         name: str,
-        instrument: InstrumentBase | None = None,
+        instrument: _InstrumentType_co = None,
         label: str | None = None,
         unit: str | None = None,
         get_cmd: str | Callable[..., Any] | Literal[False] | None = None,

src/qcodes/parameters/parameter_base.py

@@ -8,9 +8,10 @@
 from contextlib import contextmanager
 from datetime import datetime
 from functools import cached_property, wraps
-from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, overload
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, overload
 
 import numpy as np
+from typing_extensions import TypeVar
 
 from qcodes.metadatable import Metadatable, MetadatableWithName
 from qcodes.parameters import ParamSpecBase
@@ -42,6 +43,13 @@
     from qcodes.instrument import InstrumentBase
     from qcodes.logger.instrument_logger import InstrumentLoggerAdapter
 
+_InstrumentType_co = TypeVar(
+    "_InstrumentType_co",
+    bound="InstrumentBase | None",
+    default="InstrumentBase | None",
+    covariant=True,
+)
+
 LOG = logging.getLogger(__name__)
 
 
@@ -109,7 +117,7 @@ def invert_val_mapping(val_mapping: Mapping[Any, Any]) -> dict[Any, Any]:
     return {v: k for k, v in val_mapping.items()}
 
 
-class ParameterBase(MetadatableWithName):
+class ParameterBase(MetadatableWithName, Generic[_InstrumentType_co]):
     """
     Shared behavior for all parameters. Not intended to be used
     directly, normally you should use ``Parameter``, ``ArrayParameter``,
@@ -212,7 +220,7 @@ class ParameterBase(MetadatableWithName):
     def __init__(
         self,
         name: str,
-        instrument: InstrumentBase | None,
+        instrument: _InstrumentType_co = None,
         snapshot_get: bool = True,
         metadata: Mapping[Any, Any] | None = None,
         step: float | None = None,
@@ -606,7 +614,7 @@ def snapshot_base(
             state["ts"] = dttime.strftime("%Y-%m-%d %H:%M:%S")
 
         for attr in set(self._meta_attrs):
-            if attr == "instrument" and self._instrument:
+            if attr == "instrument" and self._instrument is not None:
                 state.update(
                     {
                         "instrument": full_class(self._instrument),
@@ -1033,7 +1041,7 @@ def register_name(self) -> str:
         return self._register_name or self.full_name
 
     @property
-    def instrument(self) -> InstrumentBase | None:
+    def instrument(self) -> _InstrumentType_co:
         """
         Return the first instrument that this parameter is bound to.
         E.g if this is bound to a channel it will return the channel