From c53be1359cf514cb79ea194f66791c742fb6e216 Mon Sep 17 00:00:00 2001 From: "Felt, Nicholas" Date: Wed, 30 Oct 2024 16:39:42 -0700 Subject: [PATCH 1/5] 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. From a1734c38e228c87488eb48448f8d36135ea2352b Mon Sep 17 00:00:00 2001 From: "Felt, Nicholas" Date: Thu, 31 Oct 2024 10:11:04 -0700 Subject: [PATCH 2/5] 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/basic_usage.md | 11 +++++ examples/scopes/tekscope/save_screenshot.py | 23 ++++++++++ .../screen_capture_mixin.py | 25 ++++++++-- .../drivers/scopes/tekscope/tekscope.py | 46 +++++++++++++++++-- tests/sim_devices/scope/tekscopepc.yaml | 11 +++-- tests/test_scopes.py | 33 +++++++++---- 6 files changed, 127 insertions(+), 22 deletions(-) create mode 100644 examples/scopes/tekscope/save_screenshot.py 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..52895c3f --- /dev/null +++ b/examples/scopes/tekscope/save_screenshot.py @@ -0,0 +1,23 @@ +"""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: + # Get a scope + scope: MSO6B = dm.add_scope("192.168.0.1") + + # 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 example.png, 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, will create a screenshot on the device, + # 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", local_folder="./images", keep_device_file=True) 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 index a162f8b4..b99c62a6 100644 --- 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 @@ -1,10 +1,13 @@ """A mixin class providing common methods for devices that can perform screen captures.""" -import os +from __future__ import annotations from abc import ABC, abstractmethod from pathlib import Path -from typing import final, Optional, Tuple, Union +from typing import final, Optional, Tuple, TYPE_CHECKING, Union + +if TYPE_CHECKING: + import os class ScreenCaptureMixin(ABC): @@ -23,7 +26,7 @@ def valid_image_extensions(self) -> Tuple[str, ...]: """ @final - def capture_screen( + def save_screenshot( self, filename: Union[str, os.PathLike[str]], *, @@ -51,7 +54,19 @@ def capture_screen( f"valid extensions are {self.valid_image_extensions!r}" ) raise ValueError(msg) - self._capture_screen( + 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, @@ -61,7 +76,7 @@ def capture_screen( ) @abstractmethod - def _capture_screen( + def _save_screenshot( self, filename: Path, *, diff --git a/src/tm_devices/drivers/scopes/tekscope/tekscope.py b/src/tm_devices/drivers/scopes/tekscope/tekscope.py index 521c9c73..9af13229 100644 --- a/src/tm_devices/drivers/scopes/tekscope/tekscope.py +++ b/src/tm_devices/drivers/scopes/tekscope/tekscope.py @@ -665,7 +665,40 @@ def _add_or_delete_dynamic_item( f":{item_type}:LIST? returned \"{','.join(item_list)}\"", ) - def _capture_screen( + 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, *, @@ -690,12 +723,17 @@ def _capture_screen( 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()}"') + 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_folder / filename).as_posix()}"', opc=True) + self.write(f"FILESYSTEM:DELETE {device_filepath_string}", opc=True) time.sleep(0.5) # wait to ensure the file is deleted def _reboot(self) -> None: diff --git a/tests/sim_devices/scope/tekscopepc.yaml b/tests/sim_devices/scope/tekscopepc.yaml index 402e3b37..9e1a9577 100644 --- a/tests/sim_devices/scope/tekscopepc.yaml +++ b/tests/sim_devices/scope/tekscopepc.yaml @@ -33,9 +33,14 @@ 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 "./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 diff --git a/tests/test_scopes.py b/tests/test_scopes.py index 26897b9d..5de20960 100644 --- a/tests/test_scopes.py +++ b/tests/test_scopes.py @@ -493,7 +493,15 @@ def test_tekscopepc( scope.reboot() with pytest.raises(ValueError, match="Invalid image extension: '.txt', valid extensions are"): - scope.capture_screen("temp.txt") + 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") filename = pathlib.Path("temp.png") local_file = tmp_path / filename @@ -501,27 +509,32 @@ def test_tekscopepc( "pyvisa.resources.messagebased.MessageBasedResource.read_raw", mock.MagicMock(return_value=b"1234"), ): - scope.capture_screen(filename, local_folder=tmp_path) + scope.save_screenshot(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 + 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 + local_file = tmp_path / "folder" / filename 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 + 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 "{filename.as_posix()}"' in stdout - assert f'FILESYSTEM:READFILE "{filename.as_posix()}"' in stdout - assert f'FILESYSTEM:DELETE "{filename.as_posix()}"' not 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) From 6a47c0d2ece36e4c8344781db8360e2f1f0d11ae Mon Sep 17 00:00:00 2001 From: "Felt, Nicholas" Date: Thu, 31 Oct 2024 12:59:38 -0700 Subject: [PATCH 3/5] docs: Update docstrings to indicate some parameters aren't used in certain devices --- .../abstract_device_functionality/screen_capture_mixin.py | 4 ++-- src/tm_devices/drivers/scopes/tekscope/tekscope.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) 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 index b99c62a6..e913ccd5 100644 --- 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 @@ -40,8 +40,8 @@ def save_screenshot( 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. + 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. diff --git a/src/tm_devices/drivers/scopes/tekscope/tekscope.py b/src/tm_devices/drivers/scopes/tekscope/tekscope.py index 9af13229..24f6cb33 100644 --- a/src/tm_devices/drivers/scopes/tekscope/tekscope.py +++ b/src/tm_devices/drivers/scopes/tekscope/tekscope.py @@ -713,7 +713,7 @@ def _save_screenshot( 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. + 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. From 0bbffac4aaa90675ee1df190f585a6344416c4f1 Mon Sep 17 00:00:00 2001 From: "Felt, Nicholas" Date: Thu, 31 Oct 2024 15:43:11 -0700 Subject: [PATCH 4/5] refactor: Allow leaving the filename empty for saving a screenshot, it will now default to a timestamped filename --- examples/scopes/tekscope/save_screenshot.py | 7 ++++++- .../screen_capture_mixin.py | 17 ++++++++++++--- tests/test_scopes.py | 21 +++++++++++++++---- 3 files changed, 37 insertions(+), 8 deletions(-) diff --git a/examples/scopes/tekscope/save_screenshot.py b/examples/scopes/tekscope/save_screenshot.py index 52895c3f..66f27551 100644 --- a/examples/scopes/tekscope/save_screenshot.py +++ b/examples/scopes/tekscope/save_screenshot.py @@ -4,7 +4,7 @@ from tm_devices.drivers import MSO6B with DeviceManager(verbose=True) as dm: - # Get a scope + # Add a scope scope: MSO6B = dm.add_scope("192.168.0.1") # Send some commands @@ -12,6 +12,11 @@ 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, 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, 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. 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 index e913ccd5..8147c70e 100644 --- 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 @@ -3,9 +3,12 @@ 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 @@ -28,7 +31,7 @@ def valid_image_extensions(self) -> Tuple[str, ...]: @final def save_screenshot( self, - filename: Union[str, os.PathLike[str]], + filename: Optional[Union[str, os.PathLike[str]]] = None, *, colors: Optional[str] = None, view_type: Optional[str] = None, @@ -39,7 +42,8 @@ def save_screenshot( """Capture a screenshot from the device and save it locally. Args: - filename: The name of the file to save the screenshot as. + 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 "./". @@ -47,7 +51,14 @@ def save_screenshot( keep_device_file: Whether to keep the file on the device after downloading it. Defaults to False. """ - filename_path = Path(filename) + 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}, " diff --git a/tests/test_scopes.py b/tests/test_scopes.py index 5de20960..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 @@ -503,13 +505,22 @@ def test_tekscopepc( ): scope.save_screenshot("temp.png", device_folder="filename.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"), + ), 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.save_screenshot(filename, local_folder=tmp_path) + 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 @@ -517,11 +528,13 @@ def test_tekscopepc( assert f'FILESYSTEM:READFILE "./{filename.as_posix()}"' in stdout assert f'FILESYSTEM:DELETE "./{filename.as_posix()}"' in stdout - local_file = tmp_path / "folder" / filename 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, From f23bf24ecd96ef7fd3431227f44929c427c0059a Mon Sep 17 00:00:00 2001 From: "Felt, Nicholas" Date: Thu, 31 Oct 2024 16:00:58 -0700 Subject: [PATCH 5/5] docs: Update the example code to show all the input arguments --- examples/scopes/tekscope/save_screenshot.py | 19 +++++++++++++------ .../drivers/scopes/tekscope/tekscope.py | 2 +- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/examples/scopes/tekscope/save_screenshot.py b/examples/scopes/tekscope/save_screenshot.py index 66f27551..7146ff3e 100644 --- a/examples/scopes/tekscope/save_screenshot.py +++ b/examples/scopes/tekscope/save_screenshot.py @@ -5,24 +5,31 @@ with DeviceManager(verbose=True) as dm: # Add a scope - scope: MSO6B = dm.add_scope("192.168.0.1") + 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, will create a screenshot on the device, + # 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, will create a screenshot on the device, + # 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, will create a screenshot on the device, - # copy it to ./images/example.jpg on the local machine, + # 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", local_folder="./images", keep_device_file=True) + scope.save_screenshot( + "example.jpg", + colors="INVERTED", + local_folder="./images", + device_folder="./device_folder", + keep_device_file=True, + ) diff --git a/src/tm_devices/drivers/scopes/tekscope/tekscope.py b/src/tm_devices/drivers/scopes/tekscope/tekscope.py index 24f6cb33..b3aeddaf 100644 --- a/src/tm_devices/drivers/scopes/tekscope/tekscope.py +++ b/src/tm_devices/drivers/scopes/tekscope/tekscope.py @@ -734,7 +734,7 @@ def _save_screenshot( (local_folder / filename).write_bytes(data) if not keep_device_file: self.write(f"FILESYSTEM:DELETE {device_filepath_string}", opc=True) - time.sleep(0.5) # wait to ensure the file is deleted + time.sleep(0.2) # wait to ensure the file is deleted def _reboot(self) -> None: """Reboot the device."""