From c53be1359cf514cb79ea194f66791c742fb6e216 Mon Sep 17 00:00:00 2001 From: "Felt, Nicholas" Date: Wed, 30 Oct 2024 16:39:42 -0700 Subject: [PATCH] feat: Added a screen capture mixin and implemented the functionality for the TekScope family of device drivers --- CHANGELOG.md | 2 + .../abstract_device_functionality/__init__.py | 2 + .../screen_capture_mixin.py | 84 +++++++++++++++++++ .../device_control/pi_control.py | 14 +++- .../drivers/scopes/tekscope/tekscope.py | 63 ++++++++++++-- tests/sim_devices/scope/tekscopepc.yaml | 14 ++++ tests/test_scopes.py | 39 ++++++++- 7 files changed, 205 insertions(+), 13 deletions(-) create mode 100644 src/tm_devices/driver_mixins/abstract_device_functionality/screen_capture_mixin.py 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/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..a162f8b4 --- /dev/null +++ b/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. + """ 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..521c9c73 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,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") diff --git a/tests/sim_devices/scope/tekscopepc.yaml b/tests/sim_devices/scope/tekscopepc.yaml index 06742bce..402e3b37 100644 --- a/tests/sim_devices/scope/tekscopepc.yaml +++ b/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?' diff --git a/tests/test_scopes.py b/tests/test_scopes.py index 4848f30b..26897b9d 100644 --- a/tests/test_scopes.py +++ b/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.