diff --git a/CHANGELOG.md b/CHANGELOG.md index ee0da8b2..c1c60602 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/docs/basic_usage.md b/docs/basic_usage.md index 203f2a82..47d1c957 100644 --- a/docs/basic_usage.md +++ b/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. diff --git a/examples/scopes/tekscope/save_screenshot.py b/examples/scopes/tekscope/save_screenshot.py new file mode 100644 index 00000000..7146ff3e --- /dev/null +++ b/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, + ) diff --git a/src/tm_devices/driver_mixins/abstract_device_functionality/__init__.py b/src/tm_devices/driver_mixins/abstract_device_functionality/__init__.py index cfb4b3a4..568edea5 100644 --- a/src/tm_devices/driver_mixins/abstract_device_functionality/__init__.py +++ b/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", diff --git a/src/tm_devices/driver_mixins/abstract_device_functionality/screen_capture_mixin.py b/src/tm_devices/driver_mixins/abstract_device_functionality/screen_capture_mixin.py new file mode 100644 index 00000000..8147c70e --- /dev/null +++ b/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. + """ diff --git a/src/tm_devices/driver_mixins/device_control/pi_control.py b/src/tm_devices/driver_mixins/device_control/pi_control.py index 543225fd..2d3249b2 100644 --- a/src/tm_devices/driver_mixins/device_control/pi_control.py +++ b/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.""" diff --git a/src/tm_devices/drivers/scopes/tekscope/tekscope.py b/src/tm_devices/drivers/scopes/tekscope/tekscope.py index 43967e5c..b3aeddaf 100644 --- a/src/tm_devices/drivers/scopes/tekscope/tekscope.py +++ b/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") diff --git a/tests/sim_devices/scope/tekscopepc.yaml b/tests/sim_devices/scope/tekscopepc.yaml index 06742bce..9e1a9577 100644 --- a/tests/sim_devices/scope/tekscopepc.yaml +++ b/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?' diff --git a/tests/test_scopes.py b/tests/test_scopes.py index 4848f30b..0755ee54 100644 --- a/tests/test_scopes.py +++ b/tests/test_scopes.py @@ -6,12 +6,14 @@ import subprocess import sys +from datetime import datetime from typing import cast, TYPE_CHECKING from unittest import mock import pytest import pyvisa as visa +from dateutil.tz import tzlocal from packaging.version import Version from tm_devices import DeviceManager, register_additional_usbtmc_mapping @@ -471,11 +473,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 +494,63 @@ 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.save_screenshot("temp.txt") + with pytest.raises( + ValueError, match=r"Local folder path \(filename.txt\) is a file, not a directory." + ): + scope.save_screenshot("temp.png", local_folder="filename.txt") + with pytest.raises( + ValueError, match=r"Device folder path \(filename.txt\) is a file, not a directory." + ): + scope.save_screenshot("temp.png", device_folder="filename.txt") + + with mock.patch( + "pyvisa.resources.messagebased.MessageBasedResource.read_raw", + mock.MagicMock(return_value=b"1234"), + ), mock.patch( + "pyvisa.resources.messagebased.MessageBasedResource.write", + mock.MagicMock(return_value=None), + ), mock.patch( + "pyvisa.resources.messagebased.MessageBasedResource.read", + mock.MagicMock(return_value="1"), # this mocks the *OPC? query return value + ): + scope.enable_verification = False + filename = pathlib.Path( + datetime.now(tz=tzlocal()).strftime(f"%Y%m%d_%H%M%S{scope.valid_image_extensions[0]}") + ) + local_file = tmp_path / filename + scope.save_screenshot(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.enable_verification = True + filename = pathlib.Path("temp.png") + local_file = tmp_path / "folder" / filename + scope.save_screenshot( + filename, + local_folder=local_file.parent, + device_folder="./new_folder", + 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 "./new_folder/{filename.as_posix()}"' in stdout + assert f'FILESYSTEM:READFILE "./new_folder/{filename.as_posix()}"' in stdout + assert f'FILESYSTEM:DELETE "./new_folder/{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.