feat: Added a screen capture mixin and implemented the functionality …

…for the TekScope family of device drivers

CHANGELOG.md

@@ -20,6 +20,8 @@ Things to be included in the next release go here.
 
 ### Added
 
+- Added a new mixin, `ScreenCaptureMixin`, that defines methods/properties used for capturing screenshots from devices.
+- Added screen capture capabilities to the `TekScope` family of device drivers.
 - Testing/linting on Python 3.13.
 - Added the `get_errors()` method to the `Device` class to enable easy access to the current error code and messages on any device.
 - Added more details to the Architectural Overview page of the documentation as well as highlighting to the device driver diagram on the page.

src/tm_devices/driver_mixins/abstract_device_functionality/__init__.py

@@ -14,6 +14,7 @@
 from .base_source_channel import BaseSourceChannel
 from .channel_control_mixin import ChannelControlMixin
 from .licensed_mixin import LicensedMixin
+from .screen_capture_mixin import ScreenCaptureMixin
 from .signal_generator_mixin import SignalGeneratorMixin
 from .usb_drives_mixin import USBDrivesMixin
 
@@ -29,6 +30,7 @@
     "PlotMixin",
     "PowerMixin",
     "ReferenceMixin",
+    "ScreenCaptureMixin",
     "SearchMixin",
     "SignalGeneratorMixin",
     "USBDrivesMixin",

src/tm_devices/driver_mixins/abstract_device_functionality/screen_capture_mixin.py

@@ -0,0 +1,84 @@
+"""A mixin class providing common methods for devices that can perform screen captures."""
+
+import os
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import final, Optional, Tuple, Union
+
+
+class ScreenCaptureMixin(ABC):
+    """A mixin class providing common methods for devices that can perform screen captures."""
+
+    @property
+    @abstractmethod
+    def valid_image_extensions(self) -> Tuple[str, ...]:
+        """Return a tuple of valid image extensions for this device.
+
+        The extensions will be in the format '.ext', where 'ext' is the lowercase extension,
+        e.g. (".png", ".jpg").
+
+        Returns:
+            Tuple[str, ...]: A tuple of valid, lowercase image extensions for this device.
+        """
+
+    @final
+    def capture_screen(
+        self,
+        filename: Union[str, os.PathLike[str]],
+        *,
+        colors: Optional[str] = None,
+        view_type: Optional[str] = None,
+        local_folder: Union[str, os.PathLike[str]] = "./",
+        device_folder: Union[str, os.PathLike[str]] = "./",
+        keep_device_file: bool = False,
+    ) -> None:
+        """Capture a screenshot from the device and save it locally.
+
+        Args:
+            filename: The name of the file to save the screenshot as.
+            colors: The color scheme to use for the screenshot.
+            view_type: The type of view to capture.
+            local_folder: The local folder to save the screenshot in. Defaults to "./".
+            device_folder: The folder on the device to save the screenshot in. Defaults to "./".
+            keep_device_file: Whether to keep the file on the device after downloading it.
+                Defaults to False.
+        """
+        filename_path = Path(filename)
+        if filename_path.suffix.lower() not in self.valid_image_extensions:
+            msg = (
+                f"Invalid image extension: {filename_path.suffix!r}, "
+                f"valid extensions are {self.valid_image_extensions!r}"
+            )
+            raise ValueError(msg)
+        self._capture_screen(
+            filename=filename_path,
+            colors=colors,
+            view_type=view_type,
+            local_folder=Path(local_folder),
+            device_folder=Path(device_folder),
+            keep_device_file=keep_device_file,
+        )
+
+    @abstractmethod
+    def _capture_screen(
+        self,
+        filename: Path,
+        *,
+        colors: Optional[str],
+        view_type: Optional[str],
+        local_folder: Path,
+        device_folder: Path,
+        keep_device_file: bool = False,
+    ) -> None:
+        """Capture a screenshot from the device and save it locally.
+
+        Args:
+            filename: The name of the file to save the screenshot as.
+            colors: The color scheme to use for the screenshot.
+            view_type: The type of view to capture.
+            local_folder: The local folder to save the screenshot in. Defaults to "./".
+            device_folder: The folder on the device to save the screenshot in. Defaults to "./".
+            keep_device_file: Whether to keep the file on the device after downloading it.
+                Defaults to False.
+        """

src/tm_devices/driver_mixins/device_control/pi_control.py

@@ -599,9 +599,17 @@ def read(self) -> str:
         """Return the read results from the VISA resource."""
         return self._visa_resource.read()
 
-    def read_raw(self) -> bytes:
-        """Return the read_raw results from the VISA resource."""
-        return self._visa_resource.read_raw()
+    def read_raw(self, size: Optional[int] = None) -> bytes:
+        """Return the read_raw results from the VISA resource.
+
+        Args:
+            size: The chunk size to use to perform the reading. Defaults to None,
+                meaning the resource wide set value is set.
+
+        Returns:
+            The bytes read from the device.
+        """
+        return self._visa_resource.read_raw(size)
 
     def reset_visa_timeout(self) -> None:
         """Reset the VISA timeout to the default value."""

src/tm_devices/drivers/scopes/tekscope/tekscope.py

@@ -3,10 +3,12 @@
 # pylint: disable=too-many-lines
 import math
 import os
+import time
 import warnings
 
 from abc import ABC
 from dataclasses import dataclass
+from pathlib import Path
 from types import MappingProxyType
 from typing import Any, cast, Dict, List, Literal, Optional, Tuple, Type, Union
 
@@ -22,30 +24,27 @@
     MSO6BCommands,
     MSO6Commands,
 )
-from tm_devices.driver_mixins.abstract_device_functionality.analysis_object_mixins import (
+from tm_devices.driver_mixins.abstract_device_functionality import (
+    BaseAFGSourceChannel,
     BusMixin,
+    ChannelControlMixin,
     HistogramMixin,
+    LicensedMixin,
     MathMixin,
     MeasurementsMixin,
     PlotMixin,
     PowerMixin,
     ReferenceMixin,
+    ScreenCaptureMixin,
     SearchMixin,
+    USBDrivesMixin,
 )
-from tm_devices.driver_mixins.abstract_device_functionality.base_afg_source_channel import (
-    BaseAFGSourceChannel,
-)
-from tm_devices.driver_mixins.abstract_device_functionality.channel_control_mixin import (
-    ChannelControlMixin,
-)
-from tm_devices.driver_mixins.abstract_device_functionality.licensed_mixin import LicensedMixin
 from tm_devices.driver_mixins.abstract_device_functionality.signal_generator_mixin import (
     ExtendedSourceDeviceConstants,
     ParameterBounds,
     SignalGeneratorMixin,
     SourceDeviceConstants,
 )
-from tm_devices.driver_mixins.abstract_device_functionality.usb_drives_mixin import USBDrivesMixin
 from tm_devices.driver_mixins.device_control import PIControl
 from tm_devices.driver_mixins.shared_implementations._tektronix_pi_scope_mixin import (
     _TektronixPIScopeMixin,  # pyright: ignore[reportPrivateUsage]
@@ -93,6 +92,7 @@ class AbstractTekScope(  # pylint: disable=too-many-public-methods
     PowerMixin,
     USBDrivesMixin,
     ChannelControlMixin,
+    ScreenCaptureMixin,
     ABC,
 ):
     """Base TekScope scope device driver.
@@ -249,6 +249,18 @@ def usb_drives(self) -> Tuple[str, ...]:
             self.write(f":FILESystem:CWD {original_dir}")
         return tuple(usb_drives)
 
+    @property
+    def valid_image_extensions(self) -> Tuple[str, ...]:
+        """Return a tuple of valid image extensions for this device.
+
+        The extensions will be in the format '.ext', where 'ext' is the lowercase extension,
+        e.g. (".png", ".jpg").
+
+        Returns:
+            Tuple[str, ...]: A tuple of valid, lowercase image extensions for this device.
+        """
+        return ".png", ".bmp", ".jpg", ".jpeg"
+
     ################################################################################################
     # Public Methods
     ################################################################################################
@@ -653,6 +665,39 @@ def _add_or_delete_dynamic_item(
                 f":{item_type}:LIST? returned \"{','.join(item_list)}\"",
             )
 
+    def _capture_screen(
+        self,
+        filename: Path,
+        *,
+        colors: Optional[str],
+        view_type: Optional[str],  # noqa: ARG002
+        local_folder: Path,
+        device_folder: Path,
+        keep_device_file: bool = False,
+    ) -> None:
+        """Capture a screenshot from the device and save it locally.
+
+        Args:
+            filename: The name of the file to save the screenshot as.
+            colors: The color scheme to use for the screenshot.
+            view_type: The type of view to capture.
+            local_folder: The local folder to save the screenshot in. Defaults to "./".
+            device_folder: The folder on the device to save the screenshot in. Defaults to "./".
+            keep_device_file: Whether to keep the file on the device after downloading it.
+                Defaults to False.
+        """
+        if colors:
+            self.set_and_check("SAVE:IMAGE:COMPOSITION", colors)
+        else:
+            self.set_and_check("SAVE:IMAGE:COMPOSITION", "NORMAL")
+        self.write(f'SAVE:IMAGE "{(device_folder / filename).as_posix()}"', opc=True)
+        self.write(f'FILESYSTEM:READFILE "{(device_folder / filename).as_posix()}"')
+        data = self.read_raw()
+        (local_folder / filename).write_bytes(data)
+        if not keep_device_file:
+            self.write(f'FILESYSTEM:DELETE "{(device_folder / filename).as_posix()}"', opc=True)
+            time.sleep(0.5)  # wait to ensure the file is deleted
+
     def _reboot(self) -> None:
         """Reboot the device."""
         self.write(":SCOPEAPP REBOOT")

tests/sim_devices/scope/tekscopepc.yaml

@@ -33,6 +33,20 @@ devices:
         r: '1'
       - q: '*RST'
       - q: '*CLS'
+      - q: SAVE:IMAGE "temp.png"
+      - q: FILESYSTEM:READFILE "temp.png"
+      - q: FILESYSTEM:DELETE "temp.png"
+    properties:
+      save_image_composition:
+        default: NORMAL
+        getter:
+          q: SAVE:IMAGE:COMPOSITION?
+          r: '{:s}'
+        setter:
+          q: SAVE:IMAGE:COMPOSITION {:s}
+        specs:
+          type: str
+          valid: [NORMAL, INVERTED]
     error:
       status_register:
         - q: '*ESR?'

tests/test_scopes.py

@@ -471,11 +471,15 @@ def test_tekscope3k_4k(device_manager: DeviceManager, capsys: pytest.CaptureFixt
     assert scope2.total_channels == 2
 
 
-def test_tekscopepc(device_manager: DeviceManager) -> None:
+def test_tekscopepc(
+    device_manager: DeviceManager, tmp_path: pathlib.Path, capsys: pytest.CaptureFixture[str]
+) -> None:
     """Test the TekScopePC implementation.
 
     Args:
         device_manager: The DeviceManager object.
+        tmp_path: pytest temporary directory fixture.
+        capsys: The captured stdout and stderr.
     """
     scope: TekScopePC = device_manager.add_scope("TEKSCOPEPC-HOSTNAME")
     # Assert TekScopePC device was added and aliased properly
@@ -488,6 +492,39 @@ def test_tekscopepc(device_manager: DeviceManager) -> None:
     with pytest.warns(UserWarning, match="Rebooting is not supported for the TekScopePC driver."):
         scope.reboot()
 
+    with pytest.raises(ValueError, match="Invalid image extension: '.txt', valid extensions are"):
+        scope.capture_screen("temp.txt")
+
+    filename = pathlib.Path("temp.png")
+    local_file = tmp_path / filename
+    with mock.patch(
+        "pyvisa.resources.messagebased.MessageBasedResource.read_raw",
+        mock.MagicMock(return_value=b"1234"),
+    ):
+        scope.capture_screen(filename, local_folder=tmp_path)
+        assert local_file.read_bytes() == b"1234"
+        stdout = capsys.readouterr().out
+        assert "SAVE:IMAGE:COMPOSITION NORMAL" in stdout
+        assert f'SAVE:IMAGE "{filename.as_posix()}"' in stdout
+        assert f'FILESYSTEM:READFILE "{filename.as_posix()}"' in stdout
+        assert f'FILESYSTEM:DELETE "{filename.as_posix()}"' in stdout
+
+    with mock.patch(
+        "pyvisa.resources.messagebased.MessageBasedResource.read_raw",
+        mock.MagicMock(return_value=b"5678"),
+    ):
+        scope.capture_screen(
+            filename, local_folder=tmp_path, colors="INVERTED", keep_device_file=True
+        )
+        assert local_file.read_bytes() == b"5678"
+        stdout = capsys.readouterr().out
+        assert "SAVE:IMAGE:COMPOSITION INVERTED" in stdout
+        assert f'SAVE:IMAGE "{filename.as_posix()}"' in stdout
+        assert f'FILESYSTEM:READFILE "{filename.as_posix()}"' in stdout
+        assert f'FILESYSTEM:DELETE "{filename.as_posix()}"' not in stdout
+
+    scope.expect_esr(0)
+
 
 def test_tekscope2k(device_manager: DeviceManager, tmp_path: pathlib.Path) -> None:
     """Test the TekScope2k implementation.