refactor: Add code to ensure that the folder paths exist, both locall…

…y and on TekScope.

Also added an example to the basic_usage.md file.

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.

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)

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,
         *,

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:

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

tests/test_scopes.py

@@ -493,35 +493,48 @@ 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
     with mock.patch(
         "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)