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,