More type safe Infiniium driver

docs/changes/newsfragments/8457.improved_driver

@@ -0,0 +1,5 @@
+The Infiniium driver now tests explicitly for the type of parent instrument when required and
+used ``cast`` where ``pyvisa`` types the
+return of ``read_binary_values``/``query_binary_values`` as a ``Sequence[float]``
+regardless of the requested ``container``. This replaces five type checker
+suppressions.

src/qcodes/instrument/sims/Keysight_Infiniium.yaml

@@ -0,0 +1,56 @@
+spec: "1.0"
+devices:
+  KeysightInfiniium:
+    eom:
+      GPIB INSTR:
+        q: "\n"
+        r: "\n"
+    error: ERROR
+    dialogues:
+      - q: "*IDN?"
+        r: "Keysight Technologies,MSOS254A,MY00000000,06.60.00902"
+      - q: ":SYSTem:HEADer OFF"
+      - q: ":SYST:HEAD ON"
+      - q: ":SYST:HEAD OFF"
+      - q: ":WAVEFORM:FORMAT WORD"
+      - q: ":WAVEFORM:BYTEORDER LSBFirst"
+      - q: ":WAVEFORM:STREAMING ON"
+      # Capability queries issued during ``_query_capabilities``.
+      - q: ":ACQ:BAND:TESTLIMITS?"
+        r: "1,<numeric>1.0E+7:5.0E+10"
+      - q: ":ACQ:POIN:TESTLIMITS?"
+        r: "1,<numeric>16:2000000"
+      - q: ":ACQ:SRAT:TESTLIMITS?"
+        r: "1,<numeric>1.0E+3:2.0E+11"
+      # The driver reads the bandwidth, sets it to AUTO to query the sample rate
+      # limits and then restores it. Returning the max bandwidth means the
+      # driver detects AUTO and restores it with the same command.
+      - q: ":ACQ:BAND?"
+        r: "5.0E+10"
+      - q: ":ACQ:BAND AUTO"
+      # Waveform preamble. Fields used by the driver are (0 indexed):
+      # 2: points, 4: xincrement, 5: xorigin, 7: yincrement, 8: yorigin,
+      # 21: y units (1 == Volt).
+      - q: ":WAV:PRE?"
+        r: "2,2,1000,1,1.0E-10,-5.0E-08,0,1.0E-04,2.5E-02,0,DC,1.0E-07,-5.0E-08,1.0,0.0,2024-01-01,00:00:00,0,3,1,2,1"
+      - q: ":WAV:POIN?"
+        r: "1000"
+      # Function definitions read back with the response header on. func1 is an
+      # FFT, func2 is not.
+      - q: ":FUNC1?"
+        r: ":FUNC1:FFTMAGNITUDE CHAN1"
+      - q: ":FUNC2?"
+        r: ":FUNC2:ADD CHAN1,CHAN2"
+
+    properties:
+      wav_source:
+        default: "CHAN1"
+        getter:
+          q: ":WAV:SOUR?"
+          r: "{}"
+        setter:
+          q: ":WAV:SOUR {}"
+
+resources:
+  GPIB::1::INSTR:
+    device: KeysightInfiniium

src/qcodes/instrument_drivers/Keysight/Infiniium.py

@@ -3,7 +3,7 @@
 from io import BytesIO
 from os.path import splitext
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, ClassVar, Literal
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -127,16 +127,23 @@ def __init__(
         self._yincrement = 0.0
         self._unit = 0
 
+    @property
+    def root_instrument(self) -> "KeysightInfiniium":
+        root_instrument = super().root_instrument
+        if not isinstance(root_instrument, KeysightInfiniium):
+            raise RuntimeError(
+                f"Trace parameter is not bound to a KeysightInfiniium instrument but {type(root_instrument)}"
+            )
+        return root_instrument
+
     @property
     def setpoints(self) -> "Sequence[ParameterBase]":
         """
         Overwrite setpoint parameter to update setpoints if auto_digitize is true
         """
         instrument = self.instrument
         if isinstance(instrument, KeysightInfiniiumChannel):
-            root_instrument: KeysightInfiniium
-            root_instrument = self.root_instrument  # type: ignore[assignment]
-            cache_setpoints = root_instrument.cache_setpoints()
+            cache_setpoints = self.root_instrument.cache_setpoints()
             if not cache_setpoints:
                 self.update_setpoints()
             return (instrument.time_axis,)
@@ -201,7 +208,12 @@ def update_fft_setpoints(self) -> None:
         """
         Update waveform parameters for an FFT.
         """
-        instrument: KeysightInfiniiumFunction = self.instrument  # type: ignore[assignment]
+        # only reached for a function parameter, see the caller in ``setpoints``
+        instrument = self.instrument
+        if not isinstance(instrument, KeysightInfiniiumFunction):
+            raise RuntimeError(
+                "FFT setpoints can only be updated for a function parameter."
+            )
         instrument.write(f":WAV:SOUR {self._channel}")
         preamble = instrument.ask(":WAV:PRE?").strip().split(",")
         self.update_setpoints(preamble)
@@ -215,7 +227,7 @@ def get_raw(self) -> npt.NDArray:
         """
         if self.instrument is None:
             raise RuntimeError("Cannot get data without instrument")
-        root_instr: KeysightInfiniium = self.root_instrument  # type: ignore[assignment]
+        root_instr = self.root_instrument
         # Check if we can use cached trace parameters
         if not root_instr.cache_setpoints():
             self.update_setpoints()
@@ -234,13 +246,16 @@ def get_raw(self) -> npt.NDArray:
         root_instr.write(":WAV:DATA?")
         # Ignore first two bytes, which should be "#0"
         _ = root_instr.visa_handle.read_bytes(2)
-        data: npt.NDArray
-        data = root_instr.visa_handle.read_binary_values(  # type: ignore[assignment]
-            "h",
-            container=np.ndarray,
-            header_fmt="empty",
-            expect_termination=True,
-            data_points=self._points,
+        # pyvisa types the return as a Sequence[float] regardless of ``container``
+        data = cast(
+            "npt.NDArray",
+            root_instr.visa_handle.read_binary_values(
+                "h",
+                container=np.ndarray,
+                header_fmt="empty",
+                expect_termination=True,
+                data_points=self._points,
+            ),
         )
         data = data.astype(np.float64)
         data = (data * self._yincrement) + self._yoffset
@@ -1275,15 +1290,20 @@ def screenshot(
         )
         try:
             with open(img_path, "wb") as f:
-                screen_bytes = self.visa_handle.query_binary_values(
-                    f":DISPlay:DATA? {img_type.upper()[1:]}",  # without .
-                    # https://docs.python.org/3/library/struct.html#format-characters
-                    datatype="B",  # Capitcal B for unsigned byte
-                    container=bytes,
+                # pyvisa types the return as a Sequence[float] regardless of
+                # ``container``
+                screen_bytes = cast(
+                    "bytes",
+                    self.visa_handle.query_binary_values(
+                        f":DISPlay:DATA? {img_type.upper()[1:]}",  # without .
+                        # https://docs.python.org/3/library/struct.html#format-characters
+                        datatype="B",  # Capitcal B for unsigned byte
+                        container=bytes,
+                    ),
                 )
-                f.write(screen_bytes)  # type: ignore[arg-type]
+                f.write(screen_bytes)
             print(f"Screen image written to {img_path}")
-            return np.asarray(pil_open(BytesIO(screen_bytes)))  # type: ignore[arg-type]
+            return np.asarray(pil_open(BytesIO(screen_bytes)))
         except Exception as e:
             self.log.error(f"Failed to save screenshot, Error occurred: \n{e}")
             return None

tests/drivers/test_keysight_infiniium.py

@@ -0,0 +1,123 @@
+"""Tests for the Keysight Infiniium driver using a pyvisa-sim backend."""
+
+from typing import TYPE_CHECKING, cast
+
+import pytest
+
+from qcodes.instrument import Instrument
+from qcodes.instrument_drivers.Keysight.Infiniium import (
+    DSOTraceParam,
+    KeysightInfiniium,
+    KeysightInfiniiumChannel,
+)
+from qcodes.validators import Arrays
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+@pytest.fixture(name="driver")
+def _make_driver() -> "Iterator[KeysightInfiniium]":
+    driver = KeysightInfiniium(
+        "infiniium_sim",
+        address="GPIB::1::INSTR",
+        pyvisa_sim_file="Keysight_Infiniium.yaml",
+    )
+    yield driver
+    driver.close()
+
+
+@pytest.fixture(name="orphan_trace")
+def _make_orphan_trace() -> "Iterator[DSOTraceParam]":
+    """A trace parameter attached to an instrument that is not an Infiniium."""
+    instrument = Instrument("not_an_infiniium")
+    trace = DSOTraceParam(
+        name="trace",
+        instrument=cast("KeysightInfiniiumChannel", instrument),
+        channel="CHAN1",
+        vals=Arrays(shape=(10,)),
+    )
+    yield trace
+    instrument.close()
+
+
+def test_idn(driver: KeysightInfiniium) -> None:
+    assert driver.IDN() == {
+        "vendor": "Keysight Technologies",
+        "model": "MSOS254A",
+        "serial": "MY00000000",
+        "firmware": "06.60.00902",
+    }
+
+
+def test_capabilities(driver: KeysightInfiniium) -> None:
+    assert driver.min_bw == 1.0e7
+    assert driver.max_bw == 5.0e10
+    assert driver.min_pts == 16
+    assert driver.max_pts == 2_000_000
+    assert driver.min_srat == 1.0e3
+    assert driver.max_srat == 2.0e11
+
+
+def test_channel_trace_root_instrument(driver: KeysightInfiniium) -> None:
+    assert driver.ch1.trace.root_instrument is driver
+
+
+def test_function_trace_root_instrument(driver: KeysightInfiniium) -> None:
+    assert driver.func1.trace.root_instrument is driver
+
+
+def test_root_instrument_raises_for_foreign_parent(
+    orphan_trace: DSOTraceParam,
+) -> None:
+    with pytest.raises(
+        RuntimeError, match="not bound to a KeysightInfiniium instrument"
+    ):
+        _ = orphan_trace.root_instrument
+
+
+def test_setpoints_raises_for_foreign_parent(orphan_trace: DSOTraceParam) -> None:
+    with pytest.raises(RuntimeError, match="Invalid type for parent instrument"):
+        _ = orphan_trace.setpoints
+
+
+def test_channel_setpoints_updates_time_axis(driver: KeysightInfiniium) -> None:
+    assert driver.cache_setpoints() is False
+    setpoints = driver.ch1.trace.setpoints
+    assert setpoints == (driver.ch1.time_axis,)
+    # setpoints have been refreshed from the preamble
+    assert driver.ch1.time_axis.points == 1000
+    assert driver.ch1.time_axis.xorigin == -5.0e-8
+    assert driver.ch1.time_axis.xincrement == 1.0e-10
+    assert driver.ch1.trace.unit == "V"
+
+
+def test_channel_setpoints_are_cached(driver: KeysightInfiniium) -> None:
+    driver.cache_setpoints(True)
+    setpoints = driver.ch1.trace.setpoints
+    assert setpoints == (driver.ch1.time_axis,)
+    # the preamble was never queried so the axis keeps its initial values
+    assert driver.ch1.time_axis.points == 1
+
+
+def test_function_fft_setpoints(driver: KeysightInfiniium) -> None:
+    assert driver.func1.function() == "FFTMAGNITUDE"
+    setpoints = driver.func1.trace.setpoints
+    assert setpoints == (driver.func1.frequency_axis,)
+    assert driver.func1.frequency_axis.points == 1000
+    assert driver.func1.frequency_axis.xorigin == -5.0e-8
+    assert driver.func1.frequency_axis.xincrement == 1.0e-10
+
+
+def test_function_non_fft_setpoints(driver: KeysightInfiniium) -> None:
+    assert driver.func2.function() == "ADD"
+    setpoints = driver.func2.trace.setpoints
+    assert setpoints == (driver.func2.time_axis,)
+    assert driver.func2.time_axis.points == 1000
+
+
+def test_update_fft_setpoints_raises_for_channel(driver: KeysightInfiniium) -> None:
+    with pytest.raises(
+        RuntimeError, match="FFT setpoints can only be updated for a function parameter"
+    ):
+        driver.ch1.trace.update_fft_setpoints()