Merge pull request #566 from klauer/enh_screenshot

ENH: add `--screenshot` to the typhos CLI

docs/source/upcoming_release_notes/566-screenshot.rst

@@ -0,0 +1,27 @@
+566 screenshot
+#################
+
+API Changes
+-----------
+- Added ``TyphosSuite.save_screenshot`` which takes a screenshot of the entire
+  suite as-displayed.
+- Added ``TyphosSuite.save_device_screenshots`` which takes individual
+  screenshots of each device display in the suite and saves them to the
+  provided formatted filename.
+
+Features
+--------
+- Add ``typhos --screenshot filename_pattern`` to take screenshots of typhos
+  displays prior to exiting early (in combination with ``--exit-after``).
+
+Bugfixes
+--------
+- N/A
+
+Maintenance
+-----------
+- N/A
+
+Contributors
+------------
+- klauer

typhos/cli.py

@@ -1,11 +1,14 @@
 """This module defines the ``typhos`` command line utility"""
+from __future__ import annotations
+
 import argparse
 import ast
 import inspect
 import logging
 import re
 import signal
 import sys
+import types
 from typing import Optional
 
 import coloredlogs
@@ -25,6 +28,31 @@
 
 logger = logging.getLogger(__name__)
 
+
+class TyphosArguments(types.SimpleNamespace):
+    """Type hints for ``typhos`` CLI entrypoint arguments."""
+
+    devices: list[str]
+    layout: str
+    cols: int
+    display_type: str
+    scrollable: str
+    size: Optional[str]
+    hide_displays: bool
+    happi_cfg: Optional[str]
+    fake_device: bool
+    version: bool
+    verbose: bool
+    dark: bool
+    stylesheet_override: Optional[list[str]]
+    stylesheet_add: Optional[list[str]]
+    profile_modules: Optional[list[str]]
+    profile_output: Optional[str]
+    benchmark: Optional[list[str]]
+    exit_after: Optional[float]
+    screenshot_filename: Optional[str]
+
+
 # Argument Parser Setup
 parser = argparse.ArgumentParser(
     description=(
@@ -181,6 +209,16 @@
         "seconds"
     ),
 )
+parser.add_argument(
+    '--screenshot',
+    dest="screenshot_filename",
+    help=(
+        "Save screenshot(s) of all contained TyphosDeviceDisplay instances to "
+        "this filename pattern prior to exiting early. This name may contain "
+        "f-string style variables, including: suite_title, widget_title, "
+        "device, and name."
+    ),
+)
 
 
 # Append to module docs
@@ -484,7 +522,8 @@ def typhos_run(
     initial_size: Optional[str] = None,
     show_displays: bool = True,
     exit_after: Optional[float] = None,
-) -> QtWidgets.QMainWindow:
+    screenshot_filename: Optional[str] = None,
+) -> Optional[QtWidgets.QMainWindow]:
     """
     Run the central typhos part of typhos.
 
@@ -516,6 +555,11 @@ def typhos_run(
     show_displays : bool, optional
         If True (default), open all the included device displays.
         If False, do not open any of the displays.
+    screenshot_filename : str, optional
+        Save screenshot(s) of all contained TyphosDeviceDisplay instances to
+        this filename pattern prior to exiting early. This name may contain
+        f-string style variables,  including: suite_title, widget_title,
+        device, and name.
 
     Returns
     -------
@@ -534,34 +578,42 @@ def typhos_run(
             scroll_option=scroll_option,
             show_displays=show_displays,
         )
-    if suite:
-        if initial_size is not None:
-            try:
-                initial_size = QtCore.QSize(
-                    *(int(opt) for opt in initial_size.split(','))
-                )
-            except TypeError as exc:
-                raise ValueError(
-                    "Invalid --size argument. Expected a two-element pair "
-                    "of comma-separated integers, e.g. --size 1000,1000"
-                ) from exc
-
-        def exit_early():
-            logger.warning(
-                "Exiting typhos early due to --exit-after=%s CLI argument.",
-                exit_after
+
+    if suite is None:
+        logger.debug("Suite creation failure")
+        return None
+
+    if initial_size is not None:
+        try:
+            initial_size = QtCore.QSize(
+                *(int(opt) for opt in initial_size.split(','))
             )
-            sys.exit(0)
+        except TypeError as exc:
+            raise ValueError(
+                "Invalid --size argument. Expected a two-element pair "
+                "of comma-separated integers, e.g. --size 1000,1000"
+            ) from exc
+
+    def exit_early():
+        logger.warning(
+            "Exiting typhos early due to --exit-after=%s CLI argument.",
+            exit_after
+        )
+
+        if screenshot_filename is not None:
+            suite.save_device_screenshots(screenshot_filename)
+
+        sys.exit(0)
 
-        if exit_after is not None and exit_after >= 0:
-            QtCore.QTimer.singleShot(exit_after * 1000.0, exit_early)
+    if exit_after is not None and exit_after >= 0:
+        QtCore.QTimer.singleShot(exit_after * 1000.0, exit_early)
 
-        return launch_suite(suite, initial_size=initial_size)
+    return launch_suite(suite, initial_size=initial_size)
 
 
 def typhos_cli(args):
     """Command Line Application for Typhos."""
-    args = parser.parse_args(args)
+    args = parser.parse_args(args, TyphosArguments())
 
     if args.version:
         print(f'Typhos: Version {typhos.__version__} from {typhos.__file__}')
@@ -602,6 +654,7 @@ def typhos_cli(args):
                 initial_size=args.size,
                 show_displays=not args.hide_displays,
                 exit_after=args.exit_after,
+                screenshot_filename=args.screenshot_filename,
             )
 
         return suite

typhos/display.py

@@ -1405,6 +1405,9 @@ def add_device(self, device, macros=None):
         self.macros = self._build_macros_from_device(device, macros=macros)
         self.load_best_template()
 
+        if not self.windowTitle():
+            self.setWindowTitle(getattr(device, "name", ""))
+
     def search_for_templates(self):
         """Search the filesystem for device-specific templates."""
         device = self.device

typhos/suite.py

@@ -384,7 +384,7 @@ def get_subdisplay(self, display):
 
         Parameters
         ----------
-        display :str or Device
+        display : str or Device
             Name of screen or device
 
         Returns
@@ -767,6 +767,63 @@ def save(self):
     # Add the template to the docstring
     save.__doc__ += textwrap.indent('\n' + utils.saved_template, '\t\t')
 
+    def save_screenshot(
+        self,
+        filename: str,
+    ) -> bool:
+        """Save a screenshot of this widget to ``filename``."""
+
+        image = utils.take_widget_screenshot(self)
+        if image is None:
+            logger.warning("Failed to take screenshot")
+            return False
+
+        logger.info(
+            "Saving screenshot of suite titled '%s' to '%s'",
+            self.windowTitle(), filename,
+        )
+        image.save(filename)
+        return True
+
+    def save_device_screenshots(
+        self,
+        filename_format: str,
+    ) -> dict[str, str]:
+        """Save screenshot(s) of devices to ``filename_format``."""
+
+        filenames = {}
+        for device in self.devices:
+            display = self.get_subdisplay(device)
+
+            if hasattr(display, "to_image"):
+                image = display.to_image()
+            else:
+                # This is a fallback for if/when we don't have a TyphosDisplay
+                image = utils.take_widget_screenshot(display)
+
+            suite_title = self.windowTitle()
+            widget_title = display.windowTitle()
+            if image is None:
+                logger.warning(
+                    "Failed to take screenshot of device: %s in %s",
+                    device.name, suite_title,
+                )
+                continue
+
+            filename = filename_format.format(
+                suite_title=suite_title,
+                widget_title=widget_title,
+                device=device,
+                name=device.name,
+            )
+            logger.info(
+                "Saving screenshot of '%s': '%s' to '%s'",
+                suite_title, widget_title, filename,
+            )
+            image.save(filename)
+            filenames[device.name] = filename
+        return filenames
+
     def _get_sidebar(self, widget):
         items = {}
         for group in self.top_level_groups.values():