Repository: tektronix/tm_devices (language: python)

Issue:
[FEAT]: `curve_query` support for the MSO2kB
### Description

The MSO2kB driver family should get `curve_query()` support.

### Additional Information

_No response_


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3ec1c9f6..1356684a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,6 +18,10 @@ Valid subsections within a version are:
 
 Things to be included in the next release go here.
 
+### Added
+
+- Added curve query support for the MSO2KB series scopes
+
 ---
 
 ## v2.2.2 (2024-08-14)
diff --git a/src/tm_devices/drivers/pi/scopes/tekscope_2k/tekscope_2k.py b/src/tm_devices/drivers/pi/scopes/tekscope_2k/tekscope_2k.py
index d548beda..65c55738 100644
--- a/src/tm_devices/drivers/pi/scopes/tekscope_2k/tekscope_2k.py
+++ b/src/tm_devices/drivers/pi/scopes/tekscope_2k/tekscope_2k.py
@@ -1,10 +1,16 @@
 """Base TekScope2k scope device driver module."""
 
+import warnings
+
 from abc import ABC
+from typing import Any, List, Optional
 
 from tm_devices.drivers.device import family_base_class
 from tm_devices.drivers.pi.scopes.scope import Scope
 
+# noinspection PyPep8Naming
+from tm_devices.helpers import ReadOnlyCachedProperty as cached_property  # noqa: N813
+
 
 @family_base_class
 class TekScope2k(Scope, ABC):
@@ -18,10 +24,114 @@ class TekScope2k(Scope, ABC):
     # Properties
     ################################################################################################
 
+    @cached_property
+    def total_channels(self) -> int:
+        """Return the total number of channels (all types)."""
+        return int(self.model[6])
+
     ################################################################################################
     # Public Methods
     ################################################################################################
 
+    def get_available_data_sources(self) -> List[str]:
+        """Fetch the list of currently active channel names available as data sources.
+
+        Returns:
+            The list of available data sources as strings.
+        """
+        source_string = self.query("SELECT?")
+        source_string = source_string.split(":")[-1]  # Remove :SELECT: from beginning
+        source_list = source_string.split(";")
+        available_source_list: List[str] = []
+        for source_item in source_list:
+            source_name, source_status = source_item.split(" ", maxsplit=1)
+            if source_status == "1":
+                available_source_list.append(source_name)
+        return available_source_list
+
+    def curve_query(  # pylint: disable=too-many-locals
+        self,
+        channel_num: int,
+        wfm_type: str = "TimeDomain",
+        output_csv_file: Optional[str] = None,
+    ) -> List[Any]:
+        """Perform a curve query on a specific channel.
+
+        Args:
+            channel_num: The channel number to perform the curve query on.
+            wfm_type: The type of waveform to query.
+            output_csv_file: An optional file path to a csv file to save the curve query data in.
+
+        Returns:
+            The list of waveform data.
+
+        Raises:
+            AssertionError: Indicates that an invalid waveform type was passed in to the method.
+        """
+        available_source_list = self.get_available_data_sources()
+        found = False
+
+        for source_item in available_source_list:
+            # Analog
+            if wfm_type == "TimeDomain":
+                if source_item == f"CH{channel_num}":
+                    self.set_and_check(":DATA:SOURCE", f"CH{channel_num}")
+                    found = True
+            elif wfm_type == "Digital":
+                if source_item == f"D{channel_num}":
+                    self.set_and_check(":DATA:SOURCE", f"D{channel_num}")
+                    found = True
+            else:
+                msg = f"{wfm_type} is an invalid waveform type!"
+                raise AssertionError(msg)
+            if found:
+                break  # break out of loop
+        if not found:
+            warnings.warn(f"source not available for curve query: CH{channel_num}", stacklevel=2)
+            return []
+
+        self.set_and_check(":DATA:ENC", "ASCI")
+        wfm_str = self.query(":CURVE?")
+        frames = wfm_str.splitlines()[0].split(",")
+        wfm_data = [float(frame) for frame in frames]
+
+        ymult = float(self.query("WFMO:YMU?"))
+        yoff = float(self.query("WFMO:YOF?"))
+        yzero = float(self.query("WFMO:YZE?"))
+
+        data = [((i - yoff) * ymult) + yzero for i in wfm_data]
+        wfm_data = [round(i, 3) for i in data]
+
+        if output_csv_file:
+            with open(output_csv_file, "w", encoding="UTF-8") as csv_file:
+                for frame in wfm_data:
+                    csv_file.write(str(frame))
+                    csv_file.write(",")
+
+        return wfm_data  # return list of frames
+
+    def turn_channel_on(
+        self,
+        channel_str: str,
+    ) -> None:
+        """Enables the display of a specific channel.
+
+        Args:
+            channel_str: The channel to turn on.
+        """
+        self.set_and_check(f"SELECT:{channel_str}", 1)
+
+    def turn_channel_off(
+        self,
+        channel_str: str,
+    ) -> None:
+        """Disables the display of a specific channel.
+
+        Args:
+            channel_str: The channel to turn off.
+        """
+        self.set_and_check(f"SELECT:{channel_str}", 0)
+
     ################################################################################################
     # Private Methods
     ################################################################################################
diff --git a/src/tm_devices/helpers/constants_and_dataclasses.py b/src/tm_devices/helpers/constants_and_dataclasses.py
index 0d1e33d8..69a6ca27 100644
--- a/src/tm_devices/helpers/constants_and_dataclasses.py
+++ b/src/tm_devices/helpers/constants_and_dataclasses.py
@@ -585,6 +585,9 @@ def __str__(self) -> str:
         SupportedModels.MSO2.value: USBTMCConfiguration(
             vendor_id=_TEKTRONIX_USBTMC_VENDOR_ID, model_id="0x0105"
         ),
+        SupportedModels.MSO2KB.value: USBTMCConfiguration(
+            vendor_id=_TEKTRONIX_USBTMC_VENDOR_ID, model_id="0x03A4"
+        ),
         SupportedModels.MSO4.value: USBTMCConfiguration(
             vendor_id=_TEKTRONIX_USBTMC_VENDOR_ID, model_id="0x0527"
         ),


Test patch:
diff --git a/tests/sim_devices/devices.yaml b/tests/sim_devices/devices.yaml
index bcd5c0f1..f7311880 100644
--- a/tests/sim_devices/devices.yaml
+++ b/tests/sim_devices/devices.yaml
@@ -253,6 +253,9 @@ resources:
   TCPIP::MSO2KB-HOSTNAME::INSTR:
     device: mso2kb
     filename: scope/mso2kb.yaml
+  USB0::0x0699::0x03A4::SERIAL1::INSTR:
+    device: mso2kb
+    filename: scope/mso2kb.yaml
   #   dpo2k
   TCPIP::DPO2K-HOSTNAME::INSTR:
     device: dpo2k
diff --git a/tests/sim_devices/scope/mso2kb.yaml b/tests/sim_devices/scope/mso2kb.yaml
index 72237a29..31ea8551 100644
--- a/tests/sim_devices/scope/mso2kb.yaml
+++ b/tests/sim_devices/scope/mso2kb.yaml
@@ -14,7 +14,7 @@ devices:
         r: "\n"
     dialogues:
       - q: '*IDN?'
-        r: TEKTRONIX,MSO2012B,001061,CF:91.1CT FV:v1.30
+        r: TEKTRONIX,MSO2024B,001061,CF:91.1CT FV:v1.30
       - q: '*OPC?'
         r: '1'
       - q: REBOOT
@@ -31,3 +31,101 @@ devices:
         - q: :ALLev?
           default: 0,"No events to report - queue empty"
           command_error: 113,"Undefined header; Command not found; EXAMPLE_COMMAND"
+    properties:
+      sre:
+        default: 0
+        getter:
+          q: '*SRE?'
+          r: '{:d}'
+        setter:
+          q: '*SRE {:d}'
+        specs:
+          min: 0
+          max: 255
+          type: int
+      verbosity:
+        default: 0
+        getter:
+          q: :VERBose?
+          r: '{:s}'
+        setter:
+          q: :VERBose {:s}
+      channel_select:
+        default: :SELECT:BUS1 0;BUS2 0;CH1 1;CH2 0;CH3 0;CH4 0;MATH 0;REF1 0;REF2
+          0;D0 1;D1 0;D2 0;D3 0;D4 0;D5 0;D6 0;D7 0;D8 0;D9 0;D10 0;D11 0;D12 0;D13
+          0;D14 0;D15 0,CONTROL CH1
+        getter:
+          q: SELECT?
+          r: '{:s}'
+        setter:
+          q: SELECT {:s}
+      chan_1_state:
+        default: '1'
+        getter:
+          q: SELECT:CH1?
+          r: '{:s}'
+        setter:
+          q: SELECT:CH1 {:s}
+      chan_2_state:
+        default: '1'
+        getter:
+          q: SELECT:CH2?
+          r: '{:s}'
+        setter:
+          q: SELECT:CH2 {:s}
+      chan_3_state:
+        default: '1'
+        getter:
+          q: SELECT:CH3?
+          r: '{:s}'
+        setter:
+          q: SELECT:CH3 {:s}
+      chan_4_state:
+        default: '1'
+        getter:
+          q: SELECT:CH4?
+          r: '{:s}'
+        setter:
+          q: SELECT:CH4 {:s}
+      curve:
+        default: 1,0,1,0,1
+        getter:
+          q: :CURVE?
+          r: '{:s}'
+        setter:
+          q: :CURVE {:s}
+      data_source:
+        default: CH1
+        getter:
+          q: :DATA:SOURCE?
+          r: '{:s}'
+        setter:
+          q: :DATA:SOURCE {:s}
+      data_encoding:
+        default: ASCI
+        getter:
+          q: :DATA:ENC?
+          r: '{:s}'
+        setter:
+          q: :DATA:ENC {:s}
+      ymult:
+        default: '1'
+        getter:
+          q: WFMO:YMU?
+          r: '{:s}'
+        setter:
+          q: WFMO:YMU {:s}
+      yoff:
+        default: '1'
+        getter:
+          q: WFMO:YOF?
+          r: '{:s}'
+        setter:
+          q: WFMO:YOF {:s}
+      yze:
+        default: '1'
+        getter:
+          q: WFMO:YZE?
+          r: '{:s}'
+        setter:
+          q: WFMO:YZE {:s}
diff --git a/tests/test_scopes.py b/tests/test_scopes.py
index 87559578..8be98649 100644
--- a/tests/test_scopes.py
+++ b/tests/test_scopes.py
@@ -2,6 +2,7 @@
 """Test the Scopes."""
 
 import os
+import pathlib
 import subprocess
 import sys
 
@@ -14,7 +15,7 @@
 from packaging.version import Version
 
 from tm_devices import DeviceManager
-from tm_devices.drivers import MSO2, MSO5, MSO5B, MSO6, MSO70KDX, TekScopeSW
+from tm_devices.drivers import MSO2, MSO2KB, MSO5, MSO5B, MSO6, MSO70KDX, TekScopeSW
 from tm_devices.drivers.pi.scopes.tekscope.tekscope import (
     ExtendedSourceDeviceConstants,
     ParameterBounds,
@@ -482,3 +483,44 @@ def test_tekscopesw(device_manager: DeviceManager) -> None:
     assert scope.ip_address == ""
     assert not scope.total_channels
     assert scope.all_channel_names_list == ()
+
+
+def test_tekscope2k(device_manager: DeviceManager, tmp_path: pathlib.Path) -> None:
+    """Test the tekscope2k implementation.
+
+    Args:
+        device_manager: The DeviceManager object.
+        tmp_path: pytest temporary directory function
+    """
+    device_manager.remove_all_devices()
+    # Verify hostname can be determined when only the IP is provided
+    scope: MSO2KB = device_manager.add_scope(
+        "MSO2KB-SERIAL1", alias="mso2kb", connection_type="USB"
+    )
+
+    assert scope.idn_string == "TEKTRONIX,MSO2024B,001061,CF:91.1CT FV:v1.30"
+
+    assert scope.all_channel_names_list == ("CH1", "CH2", "CH3", "CH4")
+    assert scope.total_channels == 4
+
+    # Test toggle all channels on/off
+    scope.turn_all_channels_on()
+    for ch_name in scope.all_channel_names_list:
+        assert scope.query(f"SELECT:{ch_name}?") == "1"
+
+    scope.turn_all_channels_off()
+    for ch_name in scope.all_channel_names_list:
+        assert scope.query(f"SELECT:{ch_name}?") == "0"
+
+    filename = tmp_path / "temp.txt"
+    curve = scope.curve_query(1, wfm_type="TimeDomain", output_csv_file=str(filename))
+    with open(filename, encoding="utf8") as curve_csv:
+        curve_reformatted_from_file = list(map(float, curve_csv.read()[:-1].split(",")))
+        assert curve == curve_reformatted_from_file
+
+    with pytest.raises(AssertionError):
+        scope.curve_query(5, wfm_type="FreqDomain")
+
+    assert scope.curve_query(5, wfm_type="TimeDomain") == []
+
+    assert scope.curve_query(0, wfm_type="Digital") == [1, 0, 1, 0, 1]