wip unifikace

Author: wyzula-jan

bec_widgets/widgets/progress/progress_backend.py

@@ -0,0 +1,388 @@
+from __future__ import annotations
+
+import enum
+import time
+from dataclasses import dataclass
+from typing import Literal
+
+import numpy as np
+from bec_lib import messages
+from bec_lib.endpoints import MessageEndpoints
+from qtpy.QtCore import QObject, QTimer, Signal
+
+
+class ProgressSource(enum.Enum):
+    """
+    Enum to define the source of the progress.
+    """
+
+    SCAN_PROGRESS = "scan_progress"
+    DEVICE_PROGRESS = "device_progress"
+
+
+@dataclass(frozen=True)
+class ProgressSnapshot:
+    source: ProgressSource
+    value: float
+    max_value: float
+    done: bool
+    status: Literal["open", "paused", "aborted", "halted", "closed"]
+    device: str | None = None
+    scan_id: str | None = None
+    scan_number: int | None = None
+    rid: str | None = None
+    is_new_scan: bool = False
+
+
+class ProgressTask(QObject):
+    """
+    Class to store progress information.
+    Inspired by https://github.com/Textualize/rich/blob/master/rich/progress.py
+    """
+
+    def __init__(
+        self, parent: QObject | None, value: float = 0, max_value: float = 0, done: bool = False
+    ):
+        super().__init__(parent=parent)
+        self.start_time = time.monotonic()
+        self.done = done
+        self.value = value
+        self.max_value = max_value
+        self._elapsed_time = 0
+
+        self.timer = QTimer(self)
+        self.timer.timeout.connect(self.update_elapsed_time)
+        self.timer.start(1000)
+
+    def update(self, value: float, max_value: float, done: bool = False):
+        """
+        Update the progress.
+        """
+        self.max_value = max_value
+        self.done = done
+        self.value = value
+        if done:
+            self.timer.stop()
+
+    def update_elapsed_time(self):
+        """
+        Update the time estimates. This is called every second by a QTimer.
+        """
+        self._elapsed_time = max(0.0, time.monotonic() - self.start_time)
+
+    @property
+    def percentage(self) -> float:
+        """float: Get progress of task as a percentage. If a None total was set, returns 0"""
+        if not self.max_value:
+            return 0.0
+        completed = (self.value / self.max_value) * 100.0
+        completed = min(100.0, max(0.0, completed))
+        return completed
+
+    @property
+    def speed(self) -> float:
+        """Get the estimated speed in steps per second."""
+        if self._elapsed_time == 0:
+            return 0.0
+
+        return self.value / self._elapsed_time
+
+    @property
+    def frequency(self) -> float:
+        """Get the estimated frequency in steps per second."""
+        if self.speed == 0:
+            return 0.0
+        return 1 / self.speed
+
+    @property
+    def time_elapsed(self) -> str:
+        return self._format_time(int(self._elapsed_time))
+
+    @property
+    def remaining(self) -> float:
+        """Get the estimated remaining steps."""
+        if self.done:
+            return 0.0
+        remaining = self.max_value - self.value
+        return remaining
+
+    @property
+    def time_remaining(self) -> str:
+        """
+        Get the estimated remaining time in the format HH:MM:SS.
+        """
+        if self.done or not self.speed or not self.remaining:
+            return self._format_time(0)
+        estimate = int(np.round(self.remaining / self.speed))
+
+        return self._format_time(estimate)
+
+    def _format_time(self, seconds: float) -> str:
+        """
+        Format the time in seconds to a string in the format HH:MM:SS.
+        """
+        return f"{seconds // 3600:02}:{(seconds // 60) % 60:02}:{seconds % 60:02}"
+
+
+class BECProgressTracker(QObject):
+    """
+    Shared backend for BEC scan and device progress messages.
+    """
+
+    progress_started = Signal(object)
+    progress_updated = Signal(object)
+    progress_finished = Signal(object)
+    progress_cleared = Signal()
+    source_changed = Signal(object)
+
+    def __init__(self, bec_dispatcher, parent: QObject | None = None):
+        super().__init__(parent=parent)
+        self.bec_dispatcher = bec_dispatcher
+        self._progress_source: ProgressSource | None = None
+        self._progress_device: str | None = None
+        self._queue_connected = False
+        self.task: ProgressTask | None = None
+        self.scan_number: int | None = None
+        self._active_scan_id: str | None = None
+        self._active_rid: str | None = None
+
+    @property
+    def progress_source(self) -> ProgressSource | None:
+        return self._progress_source
+
+    @property
+    def progress_device(self) -> str | None:
+        return self._progress_device
+
+    @property
+    def active_scan_id(self) -> str | None:
+        return self._active_scan_id
+
+    @property
+    def active_rid(self) -> str | None:
+        return self._active_rid
+
+    def start(
+        self,
+        *,
+        source: ProgressSource | None = ProgressSource.SCAN_PROGRESS,
+        device: str | None = None,
+        connect_queue: bool = False,
+        refresh_queue: bool = False,
+    ) -> None:
+        if source is not None:
+            self.set_progress_source(source, device=device)
+        if connect_queue:
+            self.connect_to_queue()
+        if refresh_queue:
+            self.refresh_queue()
+
+    def connect_to_queue(self) -> None:
+        if self._queue_connected:
+            return
+        self.bec_dispatcher.connect_slot(self.on_queue_update, MessageEndpoints.scan_queue_status())
+        self._queue_connected = True
+
+    def refresh_queue(self) -> None:
+        connector = getattr(self.bec_dispatcher.client, "connector", None)
+        if connector is None:
+            return
+        msg = connector.get(MessageEndpoints.scan_queue_status())
+        if msg is None:
+            return
+        self.on_queue_update(msg.content, msg.metadata)
+
+    def set_progress_source(self, source: ProgressSource, device: str | None = None) -> None:
+        if source == ProgressSource.DEVICE_PROGRESS and not device:
+            return
+        if self._progress_source == source and self._progress_device == device:
+            self.source_changed.emit(self.current_snapshot(value=0, max_value=100, done=False))
+            return
+
+        self._disconnect_progress_source()
+        self._progress_source = source
+        self._progress_device = None if source == ProgressSource.SCAN_PROGRESS else device
+        self.bec_dispatcher.connect_slot(self.on_progress_update, self._progress_endpoint())
+        self.source_changed.emit(self.current_snapshot(value=0, max_value=100, done=False))
+
+    def _disconnect_progress_source(self) -> None:
+        if self._progress_source is None:
+            return
+        self.bec_dispatcher.disconnect_slot(self.on_progress_update, self._progress_endpoint())
+        self._progress_source = None
+        self._progress_device = None
+
+    def _progress_endpoint(self):
+        if self._progress_source == ProgressSource.SCAN_PROGRESS:
+            return MessageEndpoints.scan_progress()
+        return MessageEndpoints.device_progress(device=self._progress_device)
+
+    def current_snapshot(
+        self,
+        *,
+        value: float,
+        max_value: float,
+        done: bool,
+        status: Literal["open", "paused", "aborted", "halted", "closed"] = "open",
+        is_new_scan: bool = False,
+    ) -> ProgressSnapshot:
+        source = self._progress_source or ProgressSource.SCAN_PROGRESS
+        return ProgressSnapshot(
+            source=source,
+            value=value,
+            max_value=max_value,
+            done=done,
+            status=status,
+            device=self._progress_device,
+            scan_id=self._active_scan_id,
+            scan_number=self.scan_number,
+            rid=self._active_rid,
+            is_new_scan=is_new_scan,
+        )
+
+    def _start_task(self, scan_id: str | None, rid: str | None = None) -> None:
+        if self.task is not None:
+            self.task.timer.stop()
+            self.task.deleteLater()
+        self.task = ProgressTask(parent=self)
+        self._active_scan_id = scan_id
+        self._active_rid = rid
+        self.progress_started.emit(self.current_snapshot(value=0, max_value=100, done=False))
+
+    def clear_task(self, *, emit_finished: bool = True) -> None:
+        if self.task is None:
+            self._active_scan_id = None
+            self._active_rid = None
+            self.progress_cleared.emit()
+            return
+        self.task.timer.stop()
+        self.task.deleteLater()
+        self.task = None
+        self._active_scan_id = None
+        self._active_rid = None
+        self.progress_cleared.emit()
+        if emit_finished:
+            self.progress_finished.emit(self.current_snapshot(value=0, max_value=100, done=True))
+
+    def on_queue_update(self, msg_content: dict, metadata: dict):
+        queue_info = self._extract_active_queue_info(msg_content)
+        if queue_info is None:
+            self.clear_task()
+            self.set_progress_source(ProgressSource.SCAN_PROGRESS)
+            return
+
+        active_request_block = queue_info.active_request_block
+        scan_id = active_request_block.scan_id or str(active_request_block.scan_number)
+        rid = getattr(active_request_block, "RID", None) or metadata.get("RID")
+        if self.task is None or self._active_scan_id != scan_id:
+            self._start_task(scan_id, rid=rid)
+        else:
+            self._active_rid = rid or self._active_rid
+
+        self.scan_number = active_request_block.scan_number
+        self.source_changed.emit(self.current_snapshot(value=0, max_value=100, done=False))
+
+        report_instructions = active_request_block.report_instructions
+        if not report_instructions:
+            return
+
+        instruction = report_instructions[0]
+        if "scan_progress" in instruction:
+            self.set_progress_source(ProgressSource.SCAN_PROGRESS)
+        elif "device_progress" in instruction:
+            if not instruction["device_progress"]:
+                return
+            self.set_progress_source(
+                ProgressSource.DEVICE_PROGRESS, device=instruction["device_progress"][0]
+            )
+
+    def _extract_active_queue_info(self, msg_content: dict):
+        if "queue" not in msg_content:
+            return None
+        if "primary" not in msg_content["queue"]:
+            return None
+        primary_queue = msg_content.get("queue").get("primary")
+        if primary_queue is None or not isinstance(primary_queue, messages.ScanQueueStatus):
+            return None
+        primary_queue_info = primary_queue.info
+        if len(primary_queue_info) == 0:
+            return None
+        scan_info = primary_queue_info[0]
+        if scan_info is None or scan_info.active_request_block is None:
+            return None
+        if scan_info.status.lower() != "running":
+            return None
+        return scan_info
+
+    def on_progress_update(self, msg_content: dict, metadata: dict):
+        if self._progress_source is None:
+            return
+        self.process_progress_message(self._progress_source, msg_content, metadata)
+
+    def process_progress_message(
+        self,
+        source: ProgressSource,
+        msg_content: dict,
+        metadata: dict,
+        *,
+        device: str | None = None,
+    ) -> ProgressSnapshot | None:
+        done = msg_content.get("done", False)
+        value = msg_content.get("value", 0)
+        max_value = msg_content.get("max_value", 100)
+        if done:
+            value = max_value
+        status: Literal["open", "paused", "aborted", "halted", "closed"] = metadata.get(
+            "status", "open"
+        )
+        scan_id = metadata.get("scan_id") or metadata.get("RID")
+        rid = metadata.get("RID")
+        is_new_scan = False
+        previous_scan_id = self._active_scan_id
+        previous_rid = self._active_rid
+        identity_changed = (
+            (scan_id is not None and scan_id != previous_scan_id)
+            or (rid is not None and rid != previous_rid)
+            or (previous_scan_id is None and previous_rid is None)
+        )
+
+        if self.task is None:
+            self._start_task(scan_id, rid=rid)
+            is_new_scan = identity_changed
+        elif scan_id is not None and scan_id != self._active_scan_id:
+            self._start_task(scan_id, rid=rid)
+            is_new_scan = True
+        elif rid is not None and rid != self._active_rid:
+            self._start_task(scan_id or self._active_scan_id, rid=rid)
+            is_new_scan = True
+
+        if self.task is None:
+            return None
+
+        self.task.update(value, max_value, done)
+        progress_device = device or self._progress_device
+        snapshot = ProgressSnapshot(
+            source=source,
+            value=value,
+            max_value=max_value,
+            done=done,
+            status=status,
+            device=progress_device if source == ProgressSource.DEVICE_PROGRESS else None,
+            scan_id=self._active_scan_id,
+            scan_number=self.scan_number,
+            rid=self._active_rid,
+            is_new_scan=is_new_scan,
+        )
+        self.progress_updated.emit(snapshot)
+        if done:
+            self.clear_task()
+        return snapshot
+
+    def cleanup(self) -> None:
+        self.clear_task(emit_finished=False)
+        self._disconnect_progress_source()
+        if self._queue_connected:
+            self.bec_dispatcher.disconnect_slot(
+                self.on_queue_update, MessageEndpoints.scan_queue_status()
+            )
+            self._queue_connected = False