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)