Improve types in DPO driver

Author: jenshnielsen

src/qcodes/instrument_drivers/tektronix/DPO7200xx.py

@@ -232,7 +232,7 @@ def __init__(
         """
 
 
-class TektronixDPOWaveform(InstrumentChannel):
+class TektronixDPOWaveform(InstrumentChannel["TektronixDPOChannel"]):
     """
     This submodule retrieves data from waveform sources, e.g.
     channels.
@@ -246,7 +246,7 @@ class TektronixDPOWaveform(InstrumentChannel):
 
     def __init__(
         self,
-        parent: InstrumentBase,
+        parent: "TektronixDPOChannel",
         name: str,
         identifier: str,
         **kwargs: "Unpack[InstrumentBaseKWArgs]",
@@ -345,6 +345,12 @@ def __init__(
         )
         """Parameter trace"""
 
+    @property
+    def root_instrument(self) -> "TektronixDPO7000xx":
+        root_instrument = super().root_instrument
+        assert isinstance(root_instrument, TektronixDPO7000xx)
+        return root_instrument
+
     def _get_cmd(self, cmd_string: str) -> "Callable[[], str]":
         """
         Parameters defined in this submodule require the correct
@@ -363,7 +369,7 @@ def _get_trace_data(self) -> npt.NDArray:
 
         if not waveform.is_binary():
             raw_data = self.root_instrument.visa_handle.query_ascii_values(
-                "CURVE?", container=np.array
+                "CURVE?", container=np.ndarray
             )
         else:
             bytes_per_sample = waveform.bytes_per_sample()
@@ -378,7 +384,7 @@ def _get_trace_data(self) -> npt.NDArray:
                 "CURVE?",
                 datatype=data_type,
                 is_big_endian=is_big_endian,
-                container=np.array,
+                container=np.ndarray,
             )
 
         return (raw_data - self.raw_data_offset()) * self.scale() + self.offset()
@@ -430,7 +436,7 @@ def __init__(
         )
         """Parameter is_big_endian"""
 
-        self.bytes_per_sample: Parameter = self.add_parameter(
+        self.bytes_per_sample: Parameter[int, Self] = self.add_parameter(
             "bytes_per_sample",
             get_cmd="WFMOutpre:BYT_Nr?",
             set_cmd="WFMOutpre:BYT_Nr {}",
@@ -448,7 +454,7 @@ def __init__(
         """Parameter is_binary"""
 
 
-class TektronixDPOChannel(InstrumentChannel):
+class TektronixDPOChannel(InstrumentChannel[TektronixDPO7000xx]):
     """
     The main channel module for the oscilloscope. The parameters
     defined here reflect the waveforms as they are displayed on
@@ -457,7 +463,7 @@ class TektronixDPOChannel(InstrumentChannel):
 
     def __init__(
         self,
-        parent: Instrument | InstrumentChannel,
+        parent: TektronixDPO7000xx,
         name: str,
         channel_number: int,
         **kwargs: "Unpack[InstrumentBaseKWArgs]",
@@ -546,23 +552,23 @@ def set_trace_length(self, value: int) -> None:
             value: The requested number of samples in the trace
 
         """
-        if self.root_instrument.horizontal.record_length() < value:
+        if self.parent.horizontal.record_length() < value:
             raise ValueError(
                 "Cannot set a trace length which is larger than "
                 "the record length. Please switch to manual mode "
                 "and adjust the record length first"
             )
 
-        self.root_instrument.data.start_index(1)
-        self.root_instrument.data.stop_index(value)
+        self.parent.data.start_index(1)
+        self.parent.data.stop_index(value)
 
     def set_trace_time(self, value: float) -> None:
         """
         Args:
             value: The time over which a trace is desired.
 
         """
-        sample_rate = self.root_instrument.horizontal.sample_rate()
+        sample_rate = self.parent.horizontal.sample_rate()
         required_sample_count = int(sample_rate * value)
         self.set_trace_length(required_sample_count)