bec-widgets 0.53.3__py3-none-any.whl → 0.55.0__py3-none-any.whl

bec_widgets/widgets/spiral_progress_bar/spiral_progress_bar.py

@@ -0,0 +1,594 @@
+from __future__ import annotations
+
+from typing import Literal, Optional
+
+import pyqtgraph as pg
+from bec_lib.endpoints import MessageEndpoints
+from pydantic import Field, field_validator
+from pydantic_core import PydanticCustomError
+from qtpy import QtCore, QtGui
+from qtpy.QtCore import QSize, Slot
+from qtpy.QtWidgets import QSizePolicy, QWidget
+
+from bec_widgets.utils import BECConnector, Colors, ConnectionConfig, EntryValidator
+from bec_widgets.widgets.spiral_progress_bar.ring import Ring, RingConfig
+
+
+class SpiralProgressBarConfig(ConnectionConfig):
+    color_map: str | None = Field("magma", description="Color scheme for the progress bars.")
+    min_number_of_bars: int | None = Field(
+        1, description="Minimum number of progress bars to display."
+    )
+    max_number_of_bars: int | None = Field(
+        10, description="Maximum number of progress bars to display."
+    )
+    num_bars: int | None = Field(1, description="Number of progress bars to display.")
+    gap: int | None = Field(10, description="Gap between progress bars.")
+    auto_updates: bool | None = Field(
+        True, description="Enable or disable updates based on scan queue status."
+    )
+    rings: list[RingConfig] | None = Field([], description="List of ring configurations.")
+
+    @field_validator("num_bars")
+    def validate_num_bars(cls, v, values):
+        min_number_of_bars = values.data.get("min_number_of_bars", None)
+        max_number_of_bars = values.data.get("max_number_of_bars", None)
+        if min_number_of_bars is not None and max_number_of_bars is not None:
+            print(
+                f"Number of bars adjusted to be between defined min:{min_number_of_bars} and max:{max_number_of_bars} number of bars."
+            )
+            v = max(min_number_of_bars, min(v, max_number_of_bars))
+            return v
+
+    @field_validator("rings")
+    def validate_rings(cls, v, values):
+        if v is not None and v is not []:
+            num_bars = values.data.get("num_bars", None)
+            if len(v) != num_bars:
+                raise PydanticCustomError(
+                    "different number of configs",
+                    f"Length of rings configuration ({len(v)}) does not match the number of bars ({num_bars}).",
+                    {"wrong_value": len(v)},
+                )
+            indices = [ring.index for ring in v]
+            if sorted(indices) != list(range(len(indices))):
+                raise PydanticCustomError(
+                    "wrong indices",
+                    f"Indices of ring configurations must be unique and in order from 0 to num_bars {num_bars}.",
+                    {"wrong_value": indices},
+                )
+        return v
+
+    @field_validator("color_map")
+    def validate_color_map(cls, v, values):
+        if v is not None and v != "":
+            if v not in pg.colormap.listMaps():
+                raise PydanticCustomError(
+                    "unsupported colormap",
+                    f"Colormap '{v}' not found in the current installation of pyqtgraph",
+                    {"wrong_value": v},
+                )
+        return v
+
+
+class SpiralProgressBar(BECConnector, QWidget):
+    USER_ACCESS = [
+        "get_all_rpc",
+        "rpc_id",
+        "config_dict",
+        "rings",
+        "update_config",
+        "add_ring",
+        "remove_ring",
+        "set_precision",
+        "set_min_max_values",
+        "set_number_of_bars",
+        "set_value",
+        "set_colors_from_map",
+        "set_colors_directly",
+        "set_line_widths",
+        "set_gap",
+        "set_diameter",
+        "reset_diameter",
+        "enable_auto_updates",
+    ]
+
+    def __init__(
+        self,
+        parent=None,
+        config: SpiralProgressBarConfig | dict | None = None,
+        client=None,
+        gui_id: str | None = None,
+        num_bars: int | None = None,
+    ):
+        if config is None:
+            config = SpiralProgressBarConfig(widget_class=self.__class__.__name__)
+            self.config = config
+        else:
+            if isinstance(config, dict):
+                config = SpiralProgressBarConfig(**config, widget_class=self.__class__.__name__)
+            self.config = config
+        super().__init__(client=client, config=config, gui_id=gui_id)
+        QWidget.__init__(self, parent=None)
+
+        self.get_bec_shortcuts()
+        self.entry_validator = EntryValidator(self.dev)
+
+        self.RID = None
+        self.values = None
+
+        # For updating bar behaviour
+        self._auto_updates = True
+        self._rings = []
+
+        if num_bars is not None:
+            self.config.num_bars = max(
+                self.config.min_number_of_bars, min(num_bars, self.config.max_number_of_bars)
+            )
+        self.initialize_bars()
+
+        self.enable_auto_updates(self.config.auto_updates)
+
+    @property
+    def rings(self):
+        return self._rings
+
+    @rings.setter
+    def rings(self, value):
+        self._rings = value
+
+    def update_config(self, config: SpiralProgressBarConfig | dict):
+        """
+        Update the configuration of the widget.
+
+        Args:
+            config(SpiralProgressBarConfig|dict): Configuration to update.
+        """
+        if isinstance(config, dict):
+            config = SpiralProgressBarConfig(**config, widget_class=self.__class__.__name__)
+        self.config = config
+        self.clear_all()
+
+    def initialize_bars(self):
+        """
+        Initialize the progress bars.
+        """
+        start_positions = [90 * 16] * self.config.num_bars
+        directions = [-1] * self.config.num_bars
+
+        self.config.rings = [
+            RingConfig(
+                widget_class="Ring",
+                index=i,
+                start_positions=start_positions[i],
+                directions=directions[i],
+            )
+            for i in range(self.config.num_bars)
+        ]
+        self._rings = [
+            Ring(parent_progress_widget=self, config=config) for config in self.config.rings
+        ]
+
+        if self.config.color_map:
+            self.set_colors_from_map(self.config.color_map)
+
+        min_size = self._calculate_minimum_size()
+        self.setMinimumSize(min_size)
+        self.update()
+
+    def add_ring(self, **kwargs) -> Ring:
+        """
+        Add a new progress bar.
+
+        Args:
+            **kwargs: Keyword arguments for the new progress bar.
+
+        Returns:
+            Ring: Ring object.
+        """
+        if self.config.num_bars < self.config.max_number_of_bars:
+            ring = Ring(parent_progress_widget=self, **kwargs)
+            ring.config.index = self.config.num_bars
+            self.config.num_bars += 1
+            self._rings.append(ring)
+            self.config.rings.append(ring.config)
+            if self.config.color_map:
+                self.set_colors_from_map(self.config.color_map)
+            self.update()
+            return ring
+
+    def remove_ring(self, index: int):
+        """
+        Remove a progress bar by index.
+
+        Args:
+            index(int): Index of the progress bar to remove.
+        """
+        ring = self._find_ring_by_index(index)
+        ring.cleanup()
+        self._rings.remove(ring)
+        self.config.rings.remove(ring.config)
+        self.config.num_bars -= 1
+        self._reindex_rings()
+        if self.config.color_map:
+            self.set_colors_from_map(self.config.color_map)
+        self.update()
+
+    def _reindex_rings(self):
+        """
+        Reindex the progress bars.
+        """
+        for i, ring in enumerate(self._rings):
+            ring.config.index = i
+
+    def set_precision(self, precision: int, bar_index: int = None):
+        """
+        Set the precision for the progress bars. If bar_index is not provide, the precision will be set for all progress bars.
+
+        Args:
+            precision(int): Precision for the progress bars.
+            bar_index(int): Index of the progress bar to set the precision for. If provided, only a single precision can be set.
+        """
+        if bar_index is not None:
+            bar_index = self._bar_index_check(bar_index)
+            ring = self._find_ring_by_index(bar_index)
+            ring.config.precision = precision
+        else:
+            for ring in self._rings:
+                ring.config.precision = precision
+        self.update()
+
+    def set_min_max_values(
+        self,
+        min_values: int | float | list[int | float],
+        max_values: int | float | list[int | float],
+    ):
+        """
+        Set the minimum and maximum values for the progress bars.
+
+        Args:
+            min_values(int|float | list[float]): Minimum value(s) for the progress bars. If multiple progress bars are displayed, provide a list of minimum values for each progress bar.
+            max_values(int|float | list[float]): Maximum value(s) for the progress bars. If multiple progress bars are displayed, provide a list of maximum values for each progress bar.
+        """
+        if isinstance(min_values, int) or isinstance(min_values, float):
+            min_values = [min_values]
+        if isinstance(max_values, int) or isinstance(max_values, float):
+            max_values = [max_values]
+        min_values = self._adjust_list_to_bars(min_values)
+        max_values = self._adjust_list_to_bars(max_values)
+        for ring, min_value, max_value in zip(self._rings, min_values, max_values):
+            ring.set_min_max_values(min_value, max_value)
+        self.update()
+
+    def set_number_of_bars(self, num_bars: int):
+        """
+        Set the number of progress bars to display.
+
+        Args:
+            num_bars(int): Number of progress bars to display.
+        """
+        num_bars = max(
+            self.config.min_number_of_bars, min(num_bars, self.config.max_number_of_bars)
+        )
+        if num_bars != self.config.num_bars:
+            self.config.num_bars = num_bars
+            self.initialize_bars()
+
+    def set_value(self, values: int | list, ring_index: int = None):
+        """
+        Set the values for the progress bars.
+
+        Args:
+            values(int | tuple): Value(s) for the progress bars. If multiple progress bars are displayed, provide a tuple of values for each progress bar.
+            ring_index(int): Index of the progress bar to set the value for. If provided, only a single value can be set.
+
+        Examples:
+            >>> SpiralProgressBar.set_value(50)
+            >>> SpiralProgressBar.set_value([30, 40, 50]) # (outer, middle, inner)
+            >>> SpiralProgressBar.set_value(60, bar_index=1) # Set the value for the middle progress bar.
+        """
+        if ring_index is not None:
+            ring = self._find_ring_by_index(ring_index)
+            if isinstance(values, list):
+                values = values[0]
+                print(
+                    f"Warning: Only a single value can be set for a single progress bar. Using the first value in the list {values}"
+                )
+            ring.set_value(values)
+        else:
+            if isinstance(values, int):
+                values = [values]
+                values = self._adjust_list_to_bars(values)
+            for ring, value in zip(self._rings, values):
+                ring.set_value(value)
+        self.update()
+
+    def set_colors_from_map(self, colormap, color_format: Literal["RGB", "HEX"] = "RGB"):
+        """
+        Set the colors for the progress bars from a colormap.
+
+        Args:
+            colormap(str): Name of the colormap.
+            color_format(Literal["RGB","HEX"]): Format of the returned colors ('RGB', 'HEX').
+        """
+        if colormap not in pg.colormap.listMaps():
+            raise ValueError(
+                f"Colormap '{colormap}' not found in the current installation of pyqtgraph"
+            )
+        colors = Colors.golden_angle_color(colormap, self.config.num_bars, color_format)
+        self.set_colors_directly(colors)
+        self.config.color_map = colormap
+        self.update()
+
+    def set_colors_directly(self, colors: list[str | tuple] | str | tuple, bar_index: int = None):
+        """
+        Set the colors for the progress bars directly.
+
+        Args:
+            colors(list[str | tuple] | str | tuple): Color(s) for the progress bars. If multiple progress bars are displayed, provide a list of colors for each progress bar.
+            bar_index(int): Index of the progress bar to set the color for. If provided, only a single color can be set.
+        """
+        if bar_index is not None and isinstance(colors, (str, tuple)):
+            bar_index = self._bar_index_check(bar_index)
+            ring = self._find_ring_by_index(bar_index)
+            ring.set_color(colors)
+        else:
+            if isinstance(colors, (str, tuple)):
+                colors = [colors]
+            colors = self._adjust_list_to_bars(colors)
+            for ring, color in zip(self._rings, colors):
+                ring.set_color(color)
+            self.config.color_map = None
+        self.update()
+
+    def set_line_widths(self, widths: int | list[int], bar_index: int = None):
+        """
+        Set the line widths for the progress bars.
+
+        Args:
+            widths(int | list[int]): Line width(s) for the progress bars. If multiple progress bars are displayed, provide a list of line widths for each progress bar.
+            bar_index(int): Index of the progress bar to set the line width for. If provided, only a single line width can be set.
+        """
+        if bar_index is not None:
+            bar_index = self._bar_index_check(bar_index)
+            ring = self._find_ring_by_index(bar_index)
+            if isinstance(widths, list):
+                widths = widths[0]
+                print(
+                    f"Warning: Only a single line width can be set for a single progress bar. Using the first value in the list {widths}"
+                )
+            ring.set_line_width(widths)
+        else:
+            if isinstance(widths, int):
+                widths = [widths]
+            widths = self._adjust_list_to_bars(widths)
+            self.config.gap = max(widths) * 2
+            for ring, width in zip(self._rings, widths):
+                ring.set_line_width(width)
+        min_size = self._calculate_minimum_size()
+        self.setMinimumSize(min_size)
+        self.update()
+
+    def set_gap(self, gap: int):
+        """
+        Set the gap between the progress bars.
+
+        Args:
+            gap(int): Gap between the progress bars.
+        """
+        self.config.gap = gap
+        self.update()
+
+    def set_diameter(self, diameter: int):
+        """
+        Set the diameter of the widget.
+
+        Args:
+            diameter(int): Diameter of the widget.
+        """
+        size = QSize(diameter, diameter)
+        self.resize(size)
+        self.setFixedSize(size)
+
+    def _find_ring_by_index(self, index: int) -> Ring:
+        """
+        Find the ring by index.
+
+        Args:
+            index(int): Index of the ring.
+
+        Returns:
+            Ring: Ring object.
+        """
+        found_ring = None
+        for ring in self._rings:
+            if ring.config.index == index:
+                found_ring = ring
+                break
+        if found_ring is None:
+            raise ValueError(f"Ring with index {index} not found.")
+        return found_ring
+
+    def enable_auto_updates(self, enable: bool = True):
+        """
+        Enable or disable updates based on scan status. Overrides manual updates.
+        The behaviour of the whole progress bar widget will be driven by the scan queue status.
+
+        Args:
+            enable(bool): True or False.
+
+        Returns:
+            bool: True if scan segment updates are enabled.
+        """
+
+        self._auto_updates = enable
+        if enable is True:
+            self.bec_dispatcher.connect_slot(
+                self.on_scan_queue_status, MessageEndpoints.scan_queue_status()
+            )
+        else:
+            self.bec_dispatcher.disconnect_slot(
+                self.on_scan_queue_status, MessageEndpoints.scan_queue_status()
+            )
+        return self._auto_updates
+
+    @Slot(dict, dict)
+    def on_scan_queue_status(self, msg, meta):
+        primary_queue = msg.get("queue").get("primary")
+        info = primary_queue.get("info", None)
+
+        if info:
+            active_request_block = info[0].get("active_request_block", None)
+            if active_request_block:
+                report_instructions = active_request_block.get("report_instructions", None)
+                if report_instructions:
+                    instruction_type = list(report_instructions[0].keys())[0]
+                    if instruction_type == "scan_progress":
+                        if self.config.num_bars != 1:
+                            self.set_number_of_bars(1)
+                        self._hook_scan_progress(ring_index=0)
+                    elif instruction_type == "readback":
+                        devices = report_instructions[0].get("readback").get("devices")
+                        start = report_instructions[0].get("readback").get("start")
+                        end = report_instructions[0].get("readback").get("end")
+                        if self.config.num_bars != len(devices):
+                            self.set_number_of_bars(len(devices))
+                        for index, device in enumerate(devices):
+                            self._hook_readback(index, device, start[index], end[index])
+                    else:
+                        print(f"{instruction_type} not supported yet.")
+
+                    # elif instruction_type == "device_progress":
+                    #     print("hook device_progress")
+
+    def _hook_scan_progress(self, ring_index: int = None):
+        if ring_index is not None:
+            ring = self._find_ring_by_index(ring_index)
+        else:
+            ring = self._rings[0]
+
+        if ring.config.connections.slot == "on_scan_progress":
+            return
+        else:
+            ring.set_connections("on_scan_progress", MessageEndpoints.scan_progress())
+
+    def _hook_readback(self, bar_index: int, device: str, min: float | int, max: float | int):
+        ring = self._find_ring_by_index(bar_index)
+        ring.set_min_max_values(min, max)
+        endpoint = MessageEndpoints.device_readback(device)
+        ring.set_connections("on_device_readback", endpoint)
+
+    def _adjust_list_to_bars(self, items: list) -> list:
+        """
+        Utility method to adjust the list of parameters to match the number of progress bars.
+
+        Args:
+            items(list): List of parameters for the progress bars.
+
+        Returns:
+            list: List of parameters for the progress bars.
+        """
+        if items is None:
+            raise ValueError(
+                "Items cannot be None. Please provide a list for parameters for the progress bars."
+            )
+        if not isinstance(items, list):
+            items = [items]
+        if len(items) < self.config.num_bars:
+            last_item = items[-1]
+            items.extend([last_item] * (self.config.num_bars - len(items)))
+        elif len(items) > self.config.num_bars:
+            items = items[: self.config.num_bars]
+        return items
+
+    def _bar_index_check(self, bar_index: int):
+        """
+        Utility method to check if the bar index is within the range of the number of progress bars.
+
+        Args:
+            bar_index(int): Index of the progress bar to set the value for.
+        """
+        if not (0 <= bar_index < self.config.num_bars):
+            raise ValueError(
+                f"bar_index {bar_index} out of range of number of bars {self.config.num_bars}."
+            )
+        return bar_index
+
+    def paintEvent(self, event):
+        painter = QtGui.QPainter(self)
+        painter.setRenderHint(QtGui.QPainter.Antialiasing)
+        size = min(self.width(), self.height())
+        rect = QtCore.QRect(0, 0, size, size)
+        rect.adjust(
+            max(ring.config.line_width for ring in self._rings),
+            max(ring.config.line_width for ring in self._rings),
+            -max(ring.config.line_width for ring in self._rings),
+            -max(ring.config.line_width for ring in self._rings),
+        )
+
+        for i, ring in enumerate(self._rings):
+            # Background arc
+            painter.setPen(
+                QtGui.QPen(ring.background_color, ring.config.line_width, QtCore.Qt.SolidLine)
+            )
+            offset = self.config.gap * i
+            adjusted_rect = QtCore.QRect(
+                rect.left() + offset,
+                rect.top() + offset,
+                rect.width() - 2 * offset,
+                rect.height() - 2 * offset,
+            )
+            painter.drawArc(adjusted_rect, ring.config.start_position, 360 * 16)
+
+            # Foreground arc
+            pen = QtGui.QPen(ring.color, ring.config.line_width, QtCore.Qt.SolidLine)
+            pen.setCapStyle(QtCore.Qt.RoundCap)
+            painter.setPen(pen)
+            proportion = (ring.value - ring.config.min_value) / (
+                (ring.config.max_value - ring.config.min_value) + 1e-3
+            )
+            angle = int(proportion * 360 * 16 * ring.config.direction)
+            painter.drawArc(adjusted_rect, ring.start_position, angle)
+
+    def reset_diameter(self):
+        """
+        Reset the fixed size of the widget.
+        """
+        self.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Preferred)
+        self.setMinimumSize(self._calculate_minimum_size())
+        self.setMaximumSize(16777215, 16777215)
+
+    def _calculate_minimum_size(self):
+        """
+        Calculate the minimum size of the widget.
+        """
+        if not self.config.rings:
+            print("no rings to get size from setting size to 10x10")
+            return QSize(10, 10)
+        ring_widths = [self.config.rings[i].line_width for i in range(self.config.num_bars)]
+        total_width = sum(ring_widths) + self.config.gap * (self.config.num_bars - 1)
+        diameter = total_width * 2
+        if diameter < 50:
+            diameter = 50
+        return QSize(diameter, diameter)
+
+    def sizeHint(self):
+        min_size = self._calculate_minimum_size()
+        return min_size
+
+    def clear_all(self):
+        for ring in self._rings:
+            ring.cleanup()
+            del ring
+            self._rings = []
+        self.update()
+        self.initialize_bars()
+
+    def cleanup(self):
+        self.bec_dispatcher.disconnect_slot(
+            self.on_scan_queue_status, MessageEndpoints.scan_queue_status()
+        )
+        for ring in self._rings:
+            ring.cleanup()
+            del ring
+        super().cleanup()

{bec_widgets-0.53.3.dist-info → bec_widgets-0.55.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: bec_widgets
-Version: 0.53.3
+Version: 0.55.0
 Summary: BEC Widgets
 Project-URL: Bug Tracker, https://gitlab.psi.ch/bec/bec_widgets/issues
 Project-URL: Homepage, https://gitlab.psi.ch/bec/bec_widgets