From fd5af0184279400ca6d8e5d2042f31be88d180f3 Mon Sep 17 00:00:00 2001 From: wakonig_k Date: Sun, 27 Jul 2025 11:24:20 +0200 Subject: [PATCH] feat(dock area): add screenshot toolbar action --- bec_widgets/cli/client.py | 79 ++++++++++++++++++- bec_widgets/utils/bec_widget.py | 33 +++++++- .../widgets/containers/dock/dock_area.py | 7 ++ .../positioner_box/positioner_box.py | 2 +- .../positioner_box_2d/positioner_box_2d.py | 2 +- .../control/scan_control/scan_control.py | 1 + bec_widgets/widgets/plots/heatmap/heatmap.py | 1 + bec_widgets/widgets/plots/image/image.py | 1 + .../widgets/plots/motor_map/motor_map.py | 1 + .../plots/multi_waveform/multi_waveform.py | 1 + .../scatter_waveform/scatter_waveform.py | 1 + .../widgets/plots/waveform/waveform.py | 1 + tests/unit_tests/test_bec_dock.py | 61 ++++++++++++++ 13 files changed, 185 insertions(+), 6 deletions(-) diff --git a/bec_widgets/cli/client.py b/bec_widgets/cli/client.py index 2667c9e4..68aa2c4a 100644 --- a/bec_widgets/cli/client.py +++ b/bec_widgets/cli/client.py @@ -12,7 +12,7 @@ from typing import Literal, Optional from bec_lib.logger import bec_logger -from bec_widgets.cli.rpc.rpc_base import RPCBase, rpc_call +from bec_widgets.cli.rpc.rpc_base import RPCBase, rpc_call, rpc_timeout from bec_widgets.utils.bec_plugin_helper import get_all_plugin_widgets, get_plugin_client_module logger = bec_logger.logger @@ -414,6 +414,13 @@ class BECDockArea(RPCBase): dict: The state of the dock area. """ + @rpc_timeout(None) + @rpc_call + def screenshot(self, file_name: "str | None" = None): + """ + Take a screenshot of the dock area and save it to a file. + """ + @rpc_call def restore_state( self, state: "dict" = None, missing: "Literal['ignore', 'error']" = "ignore", extra="bottom" @@ -1426,6 +1433,13 @@ class Heatmap(RPCBase): Minimum decimal places for crosshair when dynamic precision is enabled. """ + @rpc_timeout(None) + @rpc_call + def screenshot(self, file_name: "str | None" = None): + """ + Take a screenshot of the dock area and save it to a file. + """ + @property @rpc_call def color_map(self) -> "str": @@ -1964,6 +1978,13 @@ class Image(RPCBase): Minimum decimal places for crosshair when dynamic precision is enabled. """ + @rpc_timeout(None) + @rpc_call + def screenshot(self, file_name: "str | None" = None): + """ + Take a screenshot of the dock area and save it to a file. + """ + @property @rpc_call def color_map(self) -> "str": @@ -2796,6 +2817,13 @@ class MotorMap(RPCBase): The font size of the legend font. """ + @rpc_timeout(None) + @rpc_call + def screenshot(self, file_name: "str | None" = None): + """ + Take a screenshot of the dock area and save it to a file. + """ + @property @rpc_call def color(self) -> "tuple": @@ -3201,6 +3229,13 @@ class MultiWaveform(RPCBase): Minimum decimal places for crosshair when dynamic precision is enabled. """ + @rpc_timeout(None) + @rpc_call + def screenshot(self, file_name: "str | None" = None): + """ + Take a screenshot of the dock area and save it to a file. + """ + @property @rpc_call def highlighted_index(self): @@ -3415,6 +3450,13 @@ class PositionerBox(RPCBase): positioner (Positioner | str) : Positioner to set, accepts str or the device """ + @rpc_timeout(None) + @rpc_call + def screenshot(self, file_name: "str | None" = None): + """ + Take a screenshot of the dock area and save it to a file. + """ + class PositionerBox2D(RPCBase): """Simple Widget to control two positioners in box form""" @@ -3437,6 +3479,13 @@ class PositionerBox2D(RPCBase): positioner (Positioner | str) : Positioner to set, accepts str or the device """ + @rpc_timeout(None) + @rpc_call + def screenshot(self, file_name: "str | None" = None): + """ + Take a screenshot of the dock area and save it to a file. + """ + class PositionerControlLine(RPCBase): """A widget that controls a single device.""" @@ -3450,6 +3499,13 @@ class PositionerControlLine(RPCBase): positioner (Positioner | str) : Positioner to set, accepts str or the device """ + @rpc_timeout(None) + @rpc_call + def screenshot(self, file_name: "str | None" = None): + """ + Take a screenshot of the dock area and save it to a file. + """ + class PositionerGroup(RPCBase): """Simple Widget to control a positioner in box form""" @@ -3908,6 +3964,13 @@ class ScanControl(RPCBase): Cleanup the BECConnector """ + @rpc_timeout(None) + @rpc_call + def screenshot(self, file_name: "str | None" = None): + """ + Take a screenshot of the dock area and save it to a file. + """ + class ScanProgressBar(RPCBase): """Widget to display a progress bar that is hooked up to the scan progress of a scan.""" @@ -4216,6 +4279,13 @@ class ScatterWaveform(RPCBase): Minimum decimal places for crosshair when dynamic precision is enabled. """ + @rpc_timeout(None) + @rpc_call + def screenshot(self, file_name: "str | None" = None): + """ + Take a screenshot of the dock area and save it to a file. + """ + @property @rpc_call def main_curve(self) -> "ScatterCurve": @@ -4833,6 +4903,13 @@ class Waveform(RPCBase): Minimum decimal places for crosshair when dynamic precision is enabled. """ + @rpc_timeout(None) + @rpc_call + def screenshot(self, file_name: "str | None" = None): + """ + Take a screenshot of the dock area and save it to a file. + """ + @property @rpc_call def curves(self) -> "list[Curve]": diff --git a/bec_widgets/utils/bec_widget.py b/bec_widgets/utils/bec_widget.py index 7482e74b..ed58aeb2 100644 --- a/bec_widgets/utils/bec_widget.py +++ b/bec_widgets/utils/bec_widget.py @@ -1,16 +1,19 @@ from __future__ import annotations +from datetime import datetime from typing import TYPE_CHECKING import darkdetect import shiboken6 from bec_lib.logger import bec_logger -from qtpy.QtCore import QObject, Slot -from qtpy.QtWidgets import QApplication +from qtpy.QtCore import QObject +from qtpy.QtWidgets import QApplication, QFileDialog, QWidget from bec_widgets.cli.rpc.rpc_register import RPCRegister from bec_widgets.utils.bec_connector import BECConnector, ConnectionConfig from bec_widgets.utils.colors import set_theme +from bec_widgets.utils.error_popups import SafeSlot +from bec_widgets.utils.rpc_decorator import rpc_timeout if TYPE_CHECKING: # pragma: no cover from bec_widgets.widgets.containers.dock import BECDock @@ -88,7 +91,7 @@ class BECWidget(BECConnector): theme = "dark" self.apply_theme(theme) - @Slot(str) + @SafeSlot(str) def apply_theme(self, theme: str): """ Apply the theme to the widget. @@ -97,6 +100,30 @@ class BECWidget(BECConnector): theme(str, optional): The theme to be applied. """ + @SafeSlot() + @SafeSlot(str) + @rpc_timeout(None) + def screenshot(self, file_name: str | None = None): + """ + Take a screenshot of the dock area and save it to a file. + """ + if not isinstance(self, QWidget): + logger.error("Cannot take screenshot of non-QWidget instance") + return + + screenshot = self.grab() + if file_name is None: + file_name, _ = QFileDialog.getSaveFileName( + self, + "Save Screenshot", + f"bec_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png", + "PNG Files (*.png);;JPEG Files (*.jpg *.jpeg);;All Files (*)", + ) + if not file_name: + return + screenshot.save(file_name) + logger.info(f"Screenshot saved to {file_name}") + def cleanup(self): """Cleanup the widget.""" with RPCRegister.delayed_broadcast(): diff --git a/bec_widgets/widgets/containers/dock/dock_area.py b/bec_widgets/widgets/containers/dock/dock_area.py index 01e6ecf9..ca6a698b 100644 --- a/bec_widgets/widgets/containers/dock/dock_area.py +++ b/bec_widgets/widgets/containers/dock/dock_area.py @@ -71,6 +71,7 @@ class BECDockArea(BECWidget, QWidget): "detach_dock", "attach_all", "save_state", + "screenshot", "restore_state", ] @@ -267,11 +268,16 @@ class BECDockArea(BECWidget, QWidget): "restore_state", MaterialIconAction(icon_name="frame_reload", tooltip="Restore Dock State", parent=self), ) + self.toolbar.components.add_safe( + "screenshot", + MaterialIconAction(icon_name="photo_camera", tooltip="Take Screenshot", parent=self), + ) bundle = ToolbarBundle("dock_actions", self.toolbar.components) bundle.add_action("attach_all") bundle.add_action("save_state") bundle.add_action("restore_state") + bundle.add_action("screenshot") self.toolbar.add_bundle(bundle) def _hook_toolbar(self): @@ -333,6 +339,7 @@ class BECDockArea(BECWidget, QWidget): self.toolbar.components.get_action("restore_state").action.triggered.connect( self.restore_state ) + self.toolbar.components.get_action("screenshot").action.triggered.connect(self.screenshot) @SafeSlot() def _create_widget_from_toolbar(self, widget_name: str) -> None: diff --git a/bec_widgets/widgets/control/device_control/positioner_box/positioner_box/positioner_box.py b/bec_widgets/widgets/control/device_control/positioner_box/positioner_box/positioner_box.py index 64b4b3f5..78cd5fa2 100644 --- a/bec_widgets/widgets/control/device_control/positioner_box/positioner_box/positioner_box.py +++ b/bec_widgets/widgets/control/device_control/positioner_box/positioner_box/positioner_box.py @@ -33,7 +33,7 @@ class PositionerBox(PositionerBoxBase): PLUGIN = True RPC = True - USER_ACCESS = ["set_positioner"] + USER_ACCESS = ["set_positioner", "screenshot"] device_changed = Signal(str, str) # Signal emitted to inform listeners about a position update position_update = Signal(float) diff --git a/bec_widgets/widgets/control/device_control/positioner_box/positioner_box_2d/positioner_box_2d.py b/bec_widgets/widgets/control/device_control/positioner_box/positioner_box_2d/positioner_box_2d.py index 9c7dd32d..37bfc90b 100644 --- a/bec_widgets/widgets/control/device_control/positioner_box/positioner_box_2d/positioner_box_2d.py +++ b/bec_widgets/widgets/control/device_control/positioner_box/positioner_box_2d/positioner_box_2d.py @@ -34,7 +34,7 @@ class PositionerBox2D(PositionerBoxBase): PLUGIN = True RPC = True - USER_ACCESS = ["set_positioner_hor", "set_positioner_ver"] + USER_ACCESS = ["set_positioner_hor", "set_positioner_ver", "screenshot"] device_changed_hor = Signal(str, str) device_changed_ver = Signal(str, str) diff --git a/bec_widgets/widgets/control/scan_control/scan_control.py b/bec_widgets/widgets/control/scan_control/scan_control.py index cf114ee5..043250e3 100644 --- a/bec_widgets/widgets/control/scan_control/scan_control.py +++ b/bec_widgets/widgets/control/scan_control/scan_control.py @@ -45,6 +45,7 @@ class ScanControl(BECWidget, QWidget): Widget to submit new scans to the queue. """ + USER_ACCESS = ["remove", "screenshot"] PLUGIN = True ICON_NAME = "tune" ARG_BOX_POSITION: int = 2 diff --git a/bec_widgets/widgets/plots/heatmap/heatmap.py b/bec_widgets/widgets/plots/heatmap/heatmap.py index e5818f3d..3173806b 100644 --- a/bec_widgets/widgets/plots/heatmap/heatmap.py +++ b/bec_widgets/widgets/plots/heatmap/heatmap.py @@ -115,6 +115,7 @@ class Heatmap(ImageBase): "auto_range_y.setter", "minimal_crosshair_precision", "minimal_crosshair_precision.setter", + "screenshot", # ImageView Specific Settings "color_map", "color_map.setter", diff --git a/bec_widgets/widgets/plots/image/image.py b/bec_widgets/widgets/plots/image/image.py index 32cc7b55..78306272 100644 --- a/bec_widgets/widgets/plots/image/image.py +++ b/bec_widgets/widgets/plots/image/image.py @@ -91,6 +91,7 @@ class Image(ImageBase): "auto_range_y.setter", "minimal_crosshair_precision", "minimal_crosshair_precision.setter", + "screenshot", # ImageView Specific Settings "color_map", "color_map.setter", diff --git a/bec_widgets/widgets/plots/motor_map/motor_map.py b/bec_widgets/widgets/plots/motor_map/motor_map.py index 19da515f..d01af730 100644 --- a/bec_widgets/widgets/plots/motor_map/motor_map.py +++ b/bec_widgets/widgets/plots/motor_map/motor_map.py @@ -128,6 +128,7 @@ class MotorMap(PlotBase): "y_log.setter", "legend_label_size", "legend_label_size.setter", + "screenshot", # motor_map specific "color", "color.setter", diff --git a/bec_widgets/widgets/plots/multi_waveform/multi_waveform.py b/bec_widgets/widgets/plots/multi_waveform/multi_waveform.py index 9c9b80be..4a891e80 100644 --- a/bec_widgets/widgets/plots/multi_waveform/multi_waveform.py +++ b/bec_widgets/widgets/plots/multi_waveform/multi_waveform.py @@ -96,6 +96,7 @@ class MultiWaveform(PlotBase): "legend_label_size.setter", "minimal_crosshair_precision", "minimal_crosshair_precision.setter", + "screenshot", # MultiWaveform Specific RPC Access "highlighted_index", "highlighted_index.setter", diff --git a/bec_widgets/widgets/plots/scatter_waveform/scatter_waveform.py b/bec_widgets/widgets/plots/scatter_waveform/scatter_waveform.py index 191f511c..9f4adc6f 100644 --- a/bec_widgets/widgets/plots/scatter_waveform/scatter_waveform.py +++ b/bec_widgets/widgets/plots/scatter_waveform/scatter_waveform.py @@ -84,6 +84,7 @@ class ScatterWaveform(PlotBase): "legend_label_size.setter", "minimal_crosshair_precision", "minimal_crosshair_precision.setter", + "screenshot", # Scatter Waveform Specific RPC Access "main_curve", "color_map", diff --git a/bec_widgets/widgets/plots/waveform/waveform.py b/bec_widgets/widgets/plots/waveform/waveform.py index 5fff58b0..db037229 100644 --- a/bec_widgets/widgets/plots/waveform/waveform.py +++ b/bec_widgets/widgets/plots/waveform/waveform.py @@ -105,6 +105,7 @@ class Waveform(PlotBase): "legend_label_size.setter", "minimal_crosshair_precision", "minimal_crosshair_precision.setter", + "screenshot", # Waveform Specific RPC Access "curves", "x_mode", diff --git a/tests/unit_tests/test_bec_dock.py b/tests/unit_tests/test_bec_dock.py index 8aacd199..2f117ae3 100644 --- a/tests/unit_tests/test_bec_dock.py +++ b/tests/unit_tests/test_bec_dock.py @@ -1,5 +1,7 @@ # pylint: disable=missing-function-docstring, missing-module-docstring, unused-import +from unittest import mock + import pytest from bec_lib.endpoints import MessageEndpoints @@ -170,3 +172,62 @@ def test_toolbar_add_utils_progress_bar(bec_dock_area): bec_dock_area.panels["ring_progress_bar_0"].widgets[0].config.widget_class == "RingProgressBar" ) + + +def test_toolbar_screenshot_action(bec_dock_area, tmpdir): + """Test the screenshot functionality from the toolbar.""" + # Create a test screenshot file path in tmpdir + screenshot_path = tmpdir.join("test_screenshot.png") + + # Mock the QFileDialog.getSaveFileName to return a test filename + with mock.patch("bec_widgets.utils.bec_widget.QFileDialog.getSaveFileName") as mock_dialog: + mock_dialog.return_value = (str(screenshot_path), "PNG Files (*.png)") + + # Mock the screenshot.save method + with mock.patch.object(bec_dock_area, "grab") as mock_grab: + mock_screenshot = mock.MagicMock() + mock_grab.return_value = mock_screenshot + + # Trigger the screenshot action + bec_dock_area.toolbar.components.get_action("screenshot").action.trigger() + + # Verify the dialog was called with correct parameters + mock_dialog.assert_called_once() + call_args = mock_dialog.call_args[0] + assert call_args[0] == bec_dock_area # parent widget + assert call_args[1] == "Save Screenshot" # dialog title + assert call_args[2].startswith("bec_") # filename starts with bec_ + assert call_args[2].endswith(".png") # filename ends with .png + assert ( + call_args[3] == "PNG Files (*.png);;JPEG Files (*.jpg *.jpeg);;All Files (*)" + ) # file filter + + # Verify grab was called + mock_grab.assert_called_once() + + # Verify save was called with the filename + mock_screenshot.save.assert_called_once_with(str(screenshot_path)) + + +def test_toolbar_screenshot_action_cancelled(bec_dock_area): + """Test the screenshot functionality when user cancels the dialog.""" + # Mock the QFileDialog.getSaveFileName to return empty filename (cancelled) + with mock.patch("bec_widgets.utils.bec_widget.QFileDialog.getSaveFileName") as mock_dialog: + mock_dialog.return_value = ("", "") + + # Mock the screenshot.save method + with mock.patch.object(bec_dock_area, "grab") as mock_grab: + mock_screenshot = mock.MagicMock() + mock_grab.return_value = mock_screenshot + + # Trigger the screenshot action + bec_dock_area.toolbar.components.get_action("screenshot").action.trigger() + + # Verify the dialog was called + mock_dialog.assert_called_once() + + # Verify grab was called (screenshot is taken before dialog) + mock_grab.assert_called_once() + + # Verify save was NOT called since dialog was cancelled + mock_screenshot.save.assert_not_called()