Add screen capture functionality for TekScope device driver family (#342

)

* feat: Added a screen capture mixin and implemented the functionality for the TekScope family of device drivers

* refactor: Add code to ensure that the folder paths exist, both locally and on TekScope.

Also added an example to the basic_usage.md file.

* docs: Update docstrings to indicate some parameters aren't used in certain devices

* refactor: Allow leaving the filename empty for saving a screenshot, it will now default to a timestamped filename

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.

docs/basic_usage.md

@@ -83,6 +83,17 @@ on CH1 of the SCOPE.
 --8<-- "examples/scopes/tekscope/generate_internal_afg_signal.py"
 ```
 
+## Save a screenshot from the device to the local machine
+
+`tm_devices` provides the ability to save a screenshot with device drivers that inherit from the
+[`ScreenCaptureMixin`][tm_devices.driver_mixins.abstract_device_functionality.screen_capture_mixin.ScreenCaptureMixin],
+and then copy that screenshot to the local machine running the Python script.
+
+```python
+# fmt: off
+--8<-- "examples/scopes/tekscope/save_screenshot.py"
+```
+
 ## Curve query saved to csv
 
 Perform a curve query and save the results to a csv file.

examples/scopes/tekscope/save_screenshot.py

@@ -0,0 +1,35 @@
+"""Save a screenshot on the device and copy it to the local machine/environment."""
+
+from tm_devices import DeviceManager
+from tm_devices.drivers import MSO6B
+
+with DeviceManager(verbose=True) as dm:
+    # Add a scope
+    scope: MSO6B = dm.add_scope("192.168.1.5")
+
+    # Send some commands
+    scope.add_new_math("MATH1", "CH1")  # add MATH1 to CH1
+    scope.turn_channel_on("CH2")  # turn on channel 2
+    scope.set_and_check(":HORIZONTAL:SCALE", 100e-9)  # adjust horizontal scale
+
+    # Save a screenshot as a timestamped file. This will create a screenshot on the device,
+    # copy it to the current working directory on the local machine,
+    # and then delete the screenshot file from the device.
+    scope.save_screenshot()
+
+    # Save a screenshot as "example.png". This will create a screenshot on the device,
+    # copy it to the current working directory on the local machine,
+    # and then delete the screenshot file from the device.
+    scope.save_screenshot("example.png")
+
+    # Save a screenshot as "example.jpg". This will create a screenshot on the device
+    # using INVERTED colors in the "./device_folder" folder,
+    # copy it to "./images/example.jpg" on the local machine,
+    # but this time the screenshot file on the device will not be deleted.
+    scope.save_screenshot(
+        "example.jpg",
+        colors="INVERTED",
+        local_folder="./images",
+        device_folder="./device_folder",
+        keep_device_file=True,
+    )

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,110 @@
+"""A mixin class providing common methods for devices that can perform screen captures."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from datetime import datetime
+from pathlib import Path
+from typing import final, Optional, Tuple, TYPE_CHECKING, Union
+
+from dateutil.tz import tzlocal
+
+if TYPE_CHECKING:
+    import os
+
+
+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 save_screenshot(
+        self,
+        filename: Optional[Union[str, os.PathLike[str]]] = None,
+        *,
+        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. Defaults to a timestamped
+                name using the first valid image extension.
+            colors: The color scheme to use for the screenshot. (Not used by all devices)
+            view_type: The type of view to capture. (Not used by all devices)
+            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 not filename:
+            filename_path = Path(
+                datetime.now(tz=tzlocal()).strftime(
+                    f"%Y%m%d_%H%M%S{self.valid_image_extensions[0]}"
+                )
+            )
+        else:
+            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)
+        local_folder_path = Path(local_folder)
+        device_folder_path = Path(device_folder)
+        if local_folder_path.is_file() or local_folder_path.suffix:
+            msg = f"Local folder path ({local_folder_path.as_posix()}) is a file, not a directory."
+            raise ValueError(msg)
+        if device_folder_path.is_file() or device_folder_path.suffix:
+            msg = (
+                f"Device folder path ({device_folder_path.as_posix()}) is a file, not a directory."
+            )
+            raise ValueError(msg)
+        if not local_folder_path.exists():
+            local_folder_path.mkdir(parents=True)
+        self._save_screenshot(
+            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 _save_screenshot(
+        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,77 @@ def _add_or_delete_dynamic_item(
                 f":{item_type}:LIST? returned \"{','.join(item_list)}\"",
             )
 
+    def _ensure_directory_exists_on_device(self, filepath: Path) -> None:
+        """Ensure that the directory of the filepath exists on the device, creating it if necessary.
+
+        Args:
+            filepath: The filepath to check.
+        """
+        with self.temporary_verbose(False):
+            original_dir = self.query(":FILESystem:CWD?")
+            # Remove the current working directory from the front of the input filepath
+            try:
+                relative_filepath = Path(filepath.relative_to(original_dir.replace('"', "")))
+            except ValueError:
+                # The input filepath is already a relative path
+                relative_filepath = filepath
+            changed_dir = False
+            try:
+                for path_part in relative_filepath.parents:  # pragma: no cover
+                    if path_part.is_file() or path_part.suffix or not path_part.name:
+                        break
+                    path_part_string = path_part.as_posix()
+                    if path_part_string not in {
+                        x.split(";")[0]
+                        for x in self.query(
+                            ":FILESystem:LDIR?", remove_quotes=True, allow_empty=True
+                        ).split(",")
+                    }:
+                        self.write(f':FILESystem:MKDir "{path_part_string}"')
+                    changed_dir = True
+                    self.write(f':FILESystem:CWD "./{path_part_string}"')
+            finally:
+                if changed_dir:
+                    self.write(f":FILESystem:CWD {original_dir}")
+
+    def _save_screenshot(
+        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. (Not used in any TekScope drivers)
+            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")
+        device_filepath = device_folder / filename
+        device_filepath_string = (
+            f'"{"./" if not device_filepath.drive else ""}{device_filepath.as_posix()}"'
+        )
+        self._ensure_directory_exists_on_device(device_filepath)
+        self.write(f"SAVE:IMAGE {device_filepath_string}", opc=True)
+        self.write(f"FILESYSTEM:READFILE {device_filepath_string}")
+        data = self.read_raw()
+        (local_folder / filename).write_bytes(data)
+        if not keep_device_file:
+            self.write(f"FILESYSTEM:DELETE {device_filepath_string}", opc=True)
+            time.sleep(0.2)  # 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,25 @@ devices:
         r: '1'
       - q: '*RST'
       - q: '*CLS'
+      - q: SAVE:IMAGE "./temp.png"
+      - q: FILESYSTEM:READFILE "./temp.png"
+      - q: FILESYSTEM:DELETE "./temp.png"
+      - q: SAVE:IMAGE "./new_folder/temp.png"
+      - q: FILESYSTEM:READFILE "./new_folder/temp.png"
+      - q: FILESYSTEM:DELETE "./new_folder/temp.png"
+      - q: :FILESystem:MKDir "new_folder"
+      - q: :FILESystem:CWD "./new_folder"
+    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?'