From 028fb5dde0ba849cb7e4c5815d4def4627ba28be Mon Sep 17 00:00:00 2001
From: "Jens H. Nielsen" <jenshnielsen@gmail.com>
Date: Sat, 1 Nov 2025 10:21:04 +0100
Subject: [PATCH] Tektronix DPO7200xx assign modules to types

---
 .../instrument_drivers/tektronix/DPO7200xx.py    | 16 ++++++++++++----
 tests/drivers/test_tektronix_dpo7200xx.py        | 12 ++++++++----
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/src/qcodes/instrument_drivers/tektronix/DPO7200xx.py b/src/qcodes/instrument_drivers/tektronix/DPO7200xx.py
index dc9d64e9cd59..57929e2c0c58 100644
--- a/src/qcodes/instrument_drivers/tektronix/DPO7200xx.py
+++ b/src/qcodes/instrument_drivers/tektronix/DPO7200xx.py
@@ -106,10 +106,14 @@ def __init__(
             self.add_submodule(measurement_name, measurement_module)
             measurement_list.append(measurement_module)
 
-        self.add_submodule("measurement", measurement_list)
-        self.add_submodule(
+        self.measurement: ChannelList[TektronixDPOMeasurement] = self.add_submodule(
+            "measurement", measurement_list
+        )
+        """Instrument module measurement"""
+        self.statistics: TektronixDPOMeasurementStatistics = self.add_submodule(
             "statistics", TektronixDPOMeasurementStatistics(self, "statistics")
         )
+        """Instrument module statistics"""
 
         channel_list = ChannelList(self, "channel", TektronixDPOChannel)
         for channel_number in range(1, self.number_of_channels + 1):
@@ -123,7 +127,10 @@ def __init__(
             self.add_submodule(channel_name, channel_module)
             channel_list.append(channel_module)
 
-        self.add_submodule("channel", channel_list)
+        self.channel: ChannelList[TektronixDPOChannel] = self.add_submodule(
+            "channel", channel_list
+        )
+        """Instrument module channel"""
 
         self.connect_message()
 
@@ -445,9 +452,10 @@ def __init__(
         super().__init__(parent, name, **kwargs)
         self._identifier = f"CH{channel_number}"
 
-        self.add_submodule(
+        self.waveform: TektronixDPOWaveform = self.add_submodule(
             "waveform", TektronixDPOWaveform(self, "waveform", self._identifier)
         )
+        """Instrument module waveform"""
 
         self.scale: Parameter = self.add_parameter(
             "scale",
diff --git a/tests/drivers/test_tektronix_dpo7200xx.py b/tests/drivers/test_tektronix_dpo7200xx.py
index 016ff6beaf64..8d6eef682d5d 100644
--- a/tests/drivers/test_tektronix_dpo7200xx.py
+++ b/tests/drivers/test_tektronix_dpo7200xx.py
@@ -1,13 +1,17 @@
 import sys
 import timeit
+from typing import TYPE_CHECKING
 
 import pytest
 
 from qcodes.instrument_drivers.tektronix.DPO7200xx import TektronixDPO7000xx
 
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
 
 @pytest.fixture(scope="function")
-def tektronix_dpo():
+def tektronix_dpo() -> "Generator[TektronixDPO7000xx, None, None]":
     """
     A six channel-per-relay instrument
     """
@@ -24,7 +28,7 @@ def tektronix_dpo():
 @pytest.mark.xfail(
     condition=sys.platform == "win32", reason="Time resolution is too low on windows"
 )
-def test_adjust_timer(tektronix_dpo) -> None:
+def test_adjust_timer(tektronix_dpo: TektronixDPO7000xx) -> None:
     """
     After adjusting the type of the measurement or the source of the
     measurement, we need wait at least 0.1 seconds
@@ -54,7 +58,7 @@ def test_adjust_timer(tektronix_dpo) -> None:
     # measurements slightly sooner then 'minimum_adjustment_time'
 
 
-def test_measurements_return_float(tektronix_dpo) -> None:
+def test_measurements_return_float(tektronix_dpo: TektronixDPO7000xx) -> None:
     amplitude = tektronix_dpo.measurement[0].amplitude()
     assert isinstance(amplitude, float)
 
@@ -62,6 +66,6 @@ def test_measurements_return_float(tektronix_dpo) -> None:
     assert isinstance(mean_amplitude, float)
 
 
-def test_measurement_sets_state(tektronix_dpo) -> None:
+def test_measurement_sets_state(tektronix_dpo: TektronixDPO7000xx) -> None:
     tektronix_dpo.measurement[1].frequency()
     assert tektronix_dpo.measurement[1].state() == 1