From 2511056557daf0b5dd78d3e85ac4befb8bf8c316 Mon Sep 17 00:00:00 2001 From: wakonig_k Date: Wed, 9 Apr 2025 12:34:15 +0200 Subject: [PATCH] feat!: add support for auto updates --- bec_widgets/cli/auto_updates.py | 464 +++++++++++------- bec_widgets/cli/client.py | 27 +- .../widgets/containers/dock/dock_area.py | 48 +- 3 files changed, 352 insertions(+), 187 deletions(-) diff --git a/bec_widgets/cli/auto_updates.py b/bec_widgets/cli/auto_updates.py index ed3ec9c0..2598e8a6 100644 --- a/bec_widgets/cli/auto_updates.py +++ b/bec_widgets/cli/auto_updates.py @@ -1,169 +1,295 @@ -# TODO autoupdate disabled -# from __future__ import annotations -# -# import threading -# from queue import Queue -# from typing import TYPE_CHECKING -# -# from pydantic import BaseModel -# -# if TYPE_CHECKING: -# from .client import BECDockArea, BECFigure -# -# -# class ScanInfo(BaseModel): -# scan_id: str -# scan_number: int -# scan_name: str -# scan_report_devices: list -# monitored_devices: list -# status: str -# model_config: dict = {"validate_assignment": True} -# -# -# class AutoUpdates: -# create_default_dock: bool = False -# enabled: bool = False -# dock_name: str = None -# -# def __init__(self, gui: BECDockArea): -# self.gui = gui -# self._default_dock = None -# self._default_fig = None -# -# def start_default_dock(self): -# """ -# Create a default dock for the auto updates. -# """ -# self.dock_name = "default_figure" -# self._default_dock = self.gui.new(self.dock_name) -# self._default_dock.new("BECFigure") -# self._default_fig = self._default_dock.elements_list[0] -# -# @staticmethod -# def get_scan_info(msg) -> ScanInfo: -# """ -# Update the script with the given data. -# """ -# info = msg.info -# status = msg.status -# scan_id = msg.scan_id -# scan_number = info.get("scan_number", 0) -# scan_name = info.get("scan_name", "Unknown") -# scan_report_devices = info.get("scan_report_devices", []) -# monitored_devices = info.get("readout_priority", {}).get("monitored", []) -# monitored_devices = [dev for dev in monitored_devices if dev not in scan_report_devices] -# return ScanInfo( -# scan_id=scan_id, -# scan_number=scan_number, -# scan_name=scan_name, -# scan_report_devices=scan_report_devices, -# monitored_devices=monitored_devices, -# status=status, -# ) -# -# def get_default_figure(self) -> BECFigure | None: -# """ -# Get the default figure from the GUI. -# """ -# return self._default_fig -# -# def do_update(self, msg): -# """ -# Run the update function if enabled. -# """ -# if not self.enabled: -# return -# if msg.status != "open": -# return -# info = self.get_scan_info(msg) -# return self.handler(info) -# -# def get_selected_device(self, monitored_devices, selected_device): -# """ -# Get the selected device for the plot. If no device is selected, the first -# device in the monitored devices list is selected. -# """ -# if selected_device: -# return selected_device -# if len(monitored_devices) > 0: -# sel_device = monitored_devices[0] -# return sel_device -# return None -# -# def handler(self, info: ScanInfo) -> None: -# """ -# Default update function. -# """ -# if info.scan_name == "line_scan" and info.scan_report_devices: -# return self.simple_line_scan(info) -# if info.scan_name == "grid_scan" and info.scan_report_devices: -# return self.simple_grid_scan(info) -# if info.scan_report_devices: -# return self.best_effort(info) -# -# def simple_line_scan(self, info: ScanInfo) -> None: -# """ -# Simple line scan. -# """ -# fig = self.get_default_figure() -# if not fig: -# return -# dev_x = info.scan_report_devices[0] -# selected_device = yield self.gui.selected_device -# dev_y = self.get_selected_device(info.monitored_devices, selected_device) -# if not dev_y: -# return -# yield fig.clear_all() -# yield fig.plot( -# x_name=dev_x, -# y_name=dev_y, -# label=f"Scan {info.scan_number} - {dev_y}", -# title=f"Scan {info.scan_number}", -# x_label=dev_x, -# y_label=dev_y, -# ) -# -# def simple_grid_scan(self, info: ScanInfo) -> None: -# """ -# Simple grid scan. -# """ -# fig = self.get_default_figure() -# if not fig: -# return -# dev_x = info.scan_report_devices[0] -# dev_y = info.scan_report_devices[1] -# selected_device = yield self.gui.selected_device -# dev_z = self.get_selected_device(info.monitored_devices, selected_device) -# yield fig.clear_all() -# yield fig.plot( -# x_name=dev_x, -# y_name=dev_y, -# z_name=dev_z, -# label=f"Scan {info.scan_number} - {dev_z}", -# title=f"Scan {info.scan_number}", -# x_label=dev_x, -# y_label=dev_y, -# ) -# -# def best_effort(self, info: ScanInfo) -> None: -# """ -# Best effort scan. -# """ -# fig = self.get_default_figure() -# if not fig: -# return -# dev_x = info.scan_report_devices[0] -# selected_device = yield self.gui.selected_device -# dev_y = self.get_selected_device(info.monitored_devices, selected_device) -# if not dev_y: -# return -# yield fig.clear_all() -# yield fig.plot( -# x_name=dev_x, -# y_name=dev_y, -# label=f"Scan {info.scan_number} - {dev_y}", -# title=f"Scan {info.scan_number}", -# x_label=dev_x, -# y_label=dev_y, -# ) +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal, overload + +from bec_lib.endpoints import MessageEndpoints +from bec_lib.logger import bec_logger +from bec_lib.messages import ScanStatusMessage + +from bec_widgets.utils.error_popups import SafeSlot + +if TYPE_CHECKING: # pragma: no cover + from bec_widgets.utils.bec_widget import BECWidget + from bec_widgets.widgets.containers.dock.dock import BECDock + from bec_widgets.widgets.containers.dock.dock_area import BECDockArea + from bec_widgets.widgets.plots.image.image import Image + from bec_widgets.widgets.plots.motor_map.motor_map import MotorMap + from bec_widgets.widgets.plots.multi_waveform.multi_waveform import MultiWaveform + from bec_widgets.widgets.plots.scatter_waveform.scatter_waveform import ScatterWaveform + from bec_widgets.widgets.plots.waveform.waveform import Waveform + + +logger = bec_logger.logger + + +class AutoUpdates: + _default_dock: BECDock + + def __init__(self, dock_area: BECDockArea): + self.dock_area = dock_area + self.bec_dispatcher = dock_area.bec_dispatcher + self._default_dock = None # type:ignore + self.current_widget: BECWidget | None = None + self.dock_name = None + self._enabled = False + + def connect(self): + """ + Establish all connections for the auto updates. + """ + self.bec_dispatcher.connect_slot(self._on_scan_status, MessageEndpoints.scan_status()) + + def disconnect(self): + """ + Disconnect all connections for the auto updates. + """ + self.bec_dispatcher.disconnect_slot( + self._on_scan_status, MessageEndpoints.scan_status() # type:ignore + ) + + @SafeSlot() + def _on_scan_status(self, content: dict, metadata: dict) -> None: + """ + Callback for scan status messages. + """ + msg = ScanStatusMessage(**content, metadata=metadata) + + match msg.status: + case "open": + self.on_scan_open(msg) + case "closed": + self.on_scan_closed(msg) + case ["aborted", "halted"]: + self.on_scan_abort(msg) + case _: + pass + + def start_default_dock(self): + """ + Create a default dock for the auto updates. + """ + self.dock_name = "update_dock" + self._default_dock = self.dock_area.new(self.dock_name) + self.current_widget = self._default_dock.new("Waveform") + + @overload + def set_dock_to_widget(self, widget: Literal["Waveform"]) -> Waveform: ... + + @overload + def set_dock_to_widget(self, widget: Literal["Image"]) -> Image: ... + + @overload + def set_dock_to_widget(self, widget: Literal["ScatterWaveform"]) -> ScatterWaveform: ... + + @overload + def set_dock_to_widget(self, widget: Literal["MotorMap"]) -> MotorMap: ... + + @overload + def set_dock_to_widget(self, widget: Literal["MultiWaveform"]) -> MultiWaveform: ... + + def set_dock_to_widget( + self, + widget: Literal["Waveform", "Image", "ScatterWaveform", "MotorMap", "MultiWaveForm"] | str, + ) -> BECWidget: + """ + Set the dock to the widget. + + Args: + widget (str): The widget to set the dock to. Must be the name of a valid widget class. + + Returns: + BECWidget: The widget that was set. + """ + if self._default_dock is None or self.current_widget is None: + logger.warning( + f"Auto Updates: No default dock found. Creating a new one with name {self.dock_name}" + ) + self.start_default_dock() + assert self.current_widget is not None + + if not self.current_widget.__class__.__name__ == widget: + self._default_dock.delete(self.current_widget.object_name) + self.current_widget = self._default_dock.new(widget) + return self.current_widget + + def get_selected_device( + self, monitored_devices, selected_device: str | None = None + ) -> str | None: + """ + Get the selected device for the plot. If no device is selected, the first + device in the monitored devices list is selected. + """ + + if selected_device is None: + selected_device = self.dock_area.selected_device + if selected_device: + return selected_device + if len(monitored_devices) > 0: + sel_device = monitored_devices[0] + return sel_device + return None + + @property + def enabled(self) -> bool: + """ + Get the enabled status of the auto updates. + """ + return self._enabled + + @enabled.setter + def enabled(self, value: bool) -> None: + """ + Set the enabled status of the auto updates. + """ + if self._enabled == value: + return + self._enabled = value + + if value: + self.connect() + self.on_start() + else: + self.disconnect() + self.on_stop() + + ######################################################################## + ################# Update Functions ##################################### + ######################################################################## + + def simple_line_scan(self, info: ScanStatusMessage) -> None: + """ + Simple line scan. + + Args: + info (ScanStatusMessage): The scan status message. + """ + + # Set the dock to the waveform widget + wf = self.set_dock_to_widget("Waveform") + + # Get the scan report devices reported by the scan + dev_x = info.scan_report_devices[0] # type:ignore + + # For the y axis, get the selected device + dev_y = self.get_selected_device(info.readout_priority["monitored"]) # type:ignore + if not dev_y: + return + + # Clear the waveform widget and plot the data + # with the scan number and device names + # as the label and title + wf.clear_all() + wf.plot( + x_name=dev_x, + y_name=dev_y, + label=f"Scan {info.scan_number} - {dev_y}", + title=f"Scan {info.scan_number}", + x_label=dev_x, + y_label=dev_y, + ) + + logger.info( + f"Auto Update [simple_line_scan]: Started plot with: x_name={dev_x}, y_name={dev_y}" + ) + + def simple_grid_scan(self, info: ScanStatusMessage) -> None: + """ + Simple grid scan. + + Args: + info (ScanStatusMessage): The scan status message. + """ + # Set the dock to the scatter waveform widget + scatter = self.set_dock_to_widget("ScatterWaveform") + + # Get the scan report devices reported by the scan + dev_x, dev_y = info.scan_report_devices[0], info.scan_report_devices[1] # type:ignore + dev_z = self.get_selected_device(info.readout_priority["monitored"]) # type:ignore + + # Clear the scatter waveform widget and plot the data + scatter.clear_all() + scatter.plot( + x_name=dev_x, y_name=dev_y, z_name=dev_z, label=f"Scan {info.scan_number} - {dev_z}" + ) + + logger.info( + f"Auto Update [simple_grid_scan]: Started plot with: x_name={dev_x}, y_name={dev_y}, z_name={dev_z}" + ) + + def best_effort(self, info: ScanStatusMessage) -> None: + """ + Best effort scan. + + Args: + info (ScanStatusMessage): The scan status message. + """ + + # If the scan report devices are empty, there is nothing we can do + if not info.scan_report_devices: + return + dev_x = info.scan_report_devices[0] # type:ignore + dev_y = self.get_selected_device(info.readout_priority["monitored"]) # type:ignore + if not dev_y: + return + + # Set the dock to the waveform widget + wf = self.set_dock_to_widget("Waveform") + + # Clear the waveform widget and plot the data + wf.clear_all() + wf.plot( + x_name=dev_x, + y_name=dev_y, + label=f"Scan {info.scan_number} - {dev_y}", + title=f"Scan {info.scan_number}", + x_label=dev_x, + y_label=dev_y, + ) + + logger.info(f"Auto Update [best_effort]: Started plot with: x_name={dev_x}, y_name={dev_y}") + + ####################################################################### + ################# GUI Callbacks ####################################### + ####################################################################### + + def on_start(self) -> None: + """ + Procedure to run when the auto updates are enabled. + """ + self.start_default_dock() + + def on_stop(self) -> None: + """ + Procedure to run when the auto updates are disabled. + """ + + def on_scan_open(self, msg: ScanStatusMessage) -> None: + """ + Procedure to run when a scan starts. + + Args: + msg (ScanStatusMessage): The scan status message. + """ + if msg.scan_name == "line_scan" and msg.scan_report_devices: + return self.simple_line_scan(msg) + if msg.scan_name == "grid_scan" and msg.scan_report_devices: + return self.simple_grid_scan(msg) + if msg.scan_report_devices: + return self.best_effort(msg) + return None + + def on_scan_closed(self, msg: ScanStatusMessage) -> None: + """ + Procedure to run when a scan ends. + + Args: + msg (ScanStatusMessage): The scan status message. + """ + + def on_scan_abort(self, msg: ScanStatusMessage) -> None: + """ + Procedure to run when a scan is aborted. + + Args: + msg (ScanStatusMessage): The scan status message. + """ diff --git a/bec_widgets/cli/client.py b/bec_widgets/cli/client.py index 32463267..dacceb42 100644 --- a/bec_widgets/cli/client.py +++ b/bec_widgets/cli/client.py @@ -363,13 +363,6 @@ class BECDockArea(RPCBase): Return all floating docks to the dock area. """ - @property - @rpc_call - def selected_device(self) -> "str": - """ - None - """ - @rpc_call def save_state(self) -> "dict": """ @@ -379,6 +372,26 @@ class BECDockArea(RPCBase): dict: The state of the dock area. """ + @property + @rpc_call + def selected_device(self) -> "str | None": + """ + Get the selected device from the auto update config. + + Returns: + str: The selected device. If no device is selected, None is returned. + """ + + @selected_device.setter + @rpc_call + def selected_device(self) -> "str | None": + """ + Get the selected device from the auto update config. + + Returns: + str: The selected device. If no device is selected, None is returned. + """ + @rpc_call def restore_state( self, state: "dict" = None, missing: "Literal['ignore', 'error']" = "ignore", extra="bottom" diff --git a/bec_widgets/widgets/containers/dock/dock_area.py b/bec_widgets/widgets/containers/dock/dock_area.py index 473ff0f9..d0c5bc41 100644 --- a/bec_widgets/widgets/containers/dock/dock_area.py +++ b/bec_widgets/widgets/containers/dock/dock_area.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Literal, Optional +from typing import TYPE_CHECKING, Literal, Optional from weakref import WeakValueDictionary from bec_lib.endpoints import MessageEndpoints @@ -37,6 +37,9 @@ from bec_widgets.widgets.services.bec_status_box.bec_status_box import BECStatus from bec_widgets.widgets.utility.logpanel.logpanel import LogPanel from bec_widgets.widgets.utility.visual.dark_mode_button.dark_mode_button import DarkModeButton +if TYPE_CHECKING: # pragma: no cover + from bec_widgets.cli.auto_updates import AutoUpdates + logger = bec_logger.logger @@ -63,8 +66,9 @@ class BECDockArea(BECWidget, QWidget): "remove", "detach_dock", "attach_all", - "selected_device", "save_state", + "selected_device", + "selected_device.setter", "restore_state", ] @@ -96,6 +100,8 @@ class BECDockArea(BECWidget, QWidget): self.layout.setSpacing(5) self.layout.setContentsMargins(0, 0, 0, 0) + self.auto_update: AutoUpdates | None = None + self._auto_update_selected_device: str | None = None self._instructions_visible = True self.dark_mode_button = DarkModeButton(parent=self, parent_id=self.gui_id, toolbar=True) @@ -253,15 +259,33 @@ class BECDockArea(BECWidget, QWidget): ) @property - def selected_device(self) -> str: - gui_id = QApplication.instance().gui_id - auto_update_config = self.client.connector.get( - MessageEndpoints.gui_auto_update_config(gui_id) - ) - try: - return auto_update_config.selected_device - except AttributeError: - return None + def selected_device(self) -> str | None: + """ + Get the selected device from the auto update config. + + Returns: + str: The selected device. If no device is selected, None is returned. + """ + return self._auto_update_selected_device + + @selected_device.setter + def selected_device(self, value: str | None) -> None: + """ + Set the selected device in the auto update config. + + Args: + value(str): The selected device. + """ + self._auto_update_selected_device = value + + def set_auto_update(self, auto_update_cls: type[AutoUpdates]) -> None: + """ + Set the auto update object for the dock area. + + Args: + auto_update(AutoUpdates): The auto update object. + """ + self.auto_update = auto_update_cls(self) @property def panels(self) -> dict[str, BECDock]: @@ -447,6 +471,8 @@ class BECDockArea(BECWidget, QWidget): """ Cleanup the dock area. """ + if self.auto_update: + self.auto_update.enabled = False self.delete_all() self.toolbar.close() self.toolbar.deleteLater()