Add AggregatedStates and state machine

bec_ipython_client/bec_ipython_client/progressbar.py

@@ -316,5 +316,7 @@ def set_finished(self, device):
         Args:
             device (str): device name
         """
+        if device not in self.devices:
+            return
         device_index = self.devices.index(device)
         self._progress.advance(self._tasks[device_index], self.NUM_STEPS)

bec_lib/bec_lib/bl_state_machine.py

@@ -0,0 +1,120 @@
+"""
+Module for managing beamline states based on configuration files.
+
+Example of the YAML configuration file:
+``` yaml
+bl_transition_states:
+  bl_state_class: AggregatedState
+  config:
+    evaluation_method: any
+    states:
+      alignment:
+        devices:
+          samx:
+            value: 0
+            abs_tol: 0.1
+            signals:
+              velocity:
+                value: 5
+                abs_tol: 0.1
+          bpm4i:
+            value: 100
+            abs_tol: 10
+          samy:
+            at: in
+            abs_tol: 0.1
+samx_within_limits:
+  bl_state_class: DeviceWithinLimitsState
+  config:
+    device: samx
+    signal: samx
+    low_limit: -1
+    high_limit: 1
+    tolerance: 0.1
+```
+
+"""
+
+from __future__ import annotations
+
+import yaml
+
+from bec_lib import bl_states
+from bec_lib.bl_state_manager import BeamlineStateManager, _state_class_for_state_type
+
+
+class BeamlineStateMachine:
+
+    def __init__(self, manager: BeamlineStateManager) -> None:
+        self._manager = manager
+        self._configs: dict[str, bl_states.BeamlineStateConfig] = {}
+
+    def load_from_config(
+        self,
+        config_path: str | None = None,
+        config_dict: dict | None = None,
+        flush: bool = True,
+        skip_existing: bool = False,
+    ) -> None:
+        """
+        Load a state configuration from a YAML file or a dictionary. If None or both are provided,
+        an error will be raised. Config must be a dictionary/YAML file with top-level beamline
+        state names and entries containing bl_state_class/config. Any previously loaded states will be
+        cleared if flush is True.
+
+        Args:
+            config_path (str | None): The path to the YAML configuration file.
+            config_dict (dict | None): A dictionary containing the configuration.
+            flush (bool): If True, existing states in the manager will be cleared before loading new ones.
+            skip_existing (bool): If True, existing states in the manager will be skipped during loading.
+        """
+        self._check_inputs(config_path=config_path, config_dict=config_dict)
+        config_dict = self._load_config(config_path=config_path, config_dict=config_dict)
+        # Check first if the config is valid before clearing the existing states
+        configs = self._parse_config(config_dict=config_dict)
+        if flush:
+            self._manager.clear_all()
+        for config in configs:
+            self._manager.add(config, skip_existing=skip_existing)
+
+    def _check_inputs(self, config_path: str | None, config_dict: dict | None) -> None:
+        if (config_path is None and config_dict is None) or (
+            config_path is not None and config_dict is not None
+        ):
+            raise ValueError("Either config_path or config_dict must be provided, but not both.")
+
+    def _load_config(self, config_path: str | None, config_dict: dict | None) -> dict:
+        if config_path:
+            with open(config_path, "r", encoding="utf-8") as f:
+                loaded_config = yaml.safe_load(f)
+            if loaded_config is None:
+                raise ValueError("Config file is empty.")
+            return loaded_config
+        if config_dict is None:
+            raise ValueError("config_dict must be provided when config_path is not set.")
+        return config_dict
+
+    def _parse_config(self, config_dict: dict) -> list[bl_states.BeamlineStateConfig]:
+        parsed_configs: list[bl_states.BeamlineStateConfig] = []
+        for state_name, entry in config_dict.items():
+            state_config_class = self._resolve_config_class(entry["bl_state_class"])
+            config = entry["config"]
+            if not isinstance(config, dict):
+                raise ValueError(f"Config for state {state_name!r} must be a dictionary.")
+            if "name" in config and config["name"] != state_name:
+                raise ValueError(
+                    f"Config name {config['name']!r} does not match top-level state name {state_name!r}."
+                )
+            config = {**config, "name": state_name}
+            parsed_configs.append(state_config_class(**config))
+        return parsed_configs
+
+    def _resolve_config_class(
+        self, bl_state_class: str | type[bl_states.BeamlineStateConfig]
+    ) -> type[bl_states.BeamlineStateConfig]:
+        if isinstance(bl_state_class, str):
+            resolved_class = _state_class_for_state_type(bl_state_class)
+        else:
+            resolved_class = bl_state_class
+        config = resolved_class.CONFIG_CLASS
+        return config

bec_lib/bec_lib/bl_state_manager.py

@@ -2,12 +2,15 @@
 
 import inspect
 import time
+from collections.abc import Mapping, Sequence
 from inspect import Parameter, Signature
-from typing import TYPE_CHECKING, TypedDict
+from typing import TYPE_CHECKING, Any, TypedDict
 
 from pydantic import BaseModel
 from rich.console import Console
 from rich.table import Table
+from rich.text import Text
+from rich.tree import Tree
 
 from bec_lib import bl_states, messages
 from bec_lib.endpoints import MessageEndpoints
@@ -73,6 +76,47 @@ class BeamlineStateGet(TypedDict):
     label: str
 
 
+def _add_parameter_to_tree(parent: Tree, name: str, value: Any) -> None:
+    """Add a configuration value to a Rich tree without interpreting it as markup."""
+    if isinstance(value, Mapping):
+        branch = parent.add(Text(str(name), style="cyan"))
+        if not value:
+            branch.add(Text("{}", style="grey70"))
+            return
+        for child_name, child_value in value.items():
+            _add_parameter_to_tree(branch, str(child_name), child_value)
+        return
+
+    if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
+        branch = parent.add(Text(str(name), style="cyan"))
+        if not value:
+            branch.add(Text("[]", style="grey70"))
+            return
+        for index, child_value in enumerate(value):
+            _add_parameter_to_tree(branch, f"[{index}]", child_value)
+        return
+
+    line = Text()
+    line.append(f"{name}: ", style="cyan")
+    line.append(repr(value), style="grey70")
+    parent.add(line)
+
+
+def _truncate_text(value: str, max_length: int) -> str:
+    """Truncate text to a deterministic maximum length."""
+    if len(value) <= max_length:
+        return value
+    return f"{value[: max_length - 1]}…"
+
+
+def _summarize_label(label: str, max_labels: int = 3, max_length: int = 64) -> str:
+    """Bound the potentially large set of matching labels shown by ``show_all``."""
+    labels = label.split("|")
+    if len(labels) > max_labels:
+        label = f"{'|'.join(labels[:max_labels])}|… (+{len(labels) - max_labels})"
+    return _truncate_text(label.replace("\n", " "), max_length=max_length)
+
+
 class BeamlineStateClientBase:
     """Base class for beamline state clients."""
 
@@ -118,6 +162,24 @@ def get(self) -> BeamlineStateGet:
         msg = msg_container["data"]
         return {"status": msg.status, "label": msg.label}
 
+    def describe(self) -> None:
+        """Pretty print the complete, validated configuration for this state."""
+        title = Text(self._state.name, style="bold magenta")
+        title.append(f" ({self._state.state_type})", style="grey70")
+        tree = Tree(title)
+        parameters = self._state.model_dump(exclude={"name"}, exclude_none=True)
+        if (
+            isinstance(self._state, bl_states.AggregatedStateConfig)
+            and self._state.evaluation_method is None
+        ):
+            parameters = {"evaluation_method": None, **parameters}
+        if not parameters:
+            tree.add(Text("No parameters", style="grey70"))
+        else:
+            for name, value in parameters.items():
+                _add_parameter_to_tree(tree, name, value)
+        Console().print(tree)
+
     def remove(self) -> None:
         """
         Remove the current beamline state.
@@ -221,17 +283,27 @@ def _wait_for_initial_state(self, state_name: str, timeout_s: float = 5.0) -> No
     ##### Public API #########
     ##########################
 
-    def add(self, state: bl_states.BeamlineStateConfig) -> None:
+    def add(self, state: bl_states.BeamlineStateConfig, skip_existing: bool = False) -> None:
         """
         Add a new beamline state to the manager.
         Args:
             state (BeamlineStateConfig): The beamline state to add.
+            skip_existing (bool): If True, existing states in the manager will be skipped during loading.
         """
-
+        if skip_existing and state.name in self._states:
+            return
         self._add_state(state)
         self._publish_states()
         self._wait_for_initial_state(state.name)
 
+    def clear_all(self) -> None:
+        """
+        Clear all beamline states from the manager.
+        """
+        for state_name in list(self._states.keys()):
+            self._delete_state(state_name)
+        self._publish_states()
+
     def delete(self, state_name: str) -> None:
         """
         Delete a beamline state from the manager.
@@ -257,12 +329,6 @@ def show_all(self):
         Pretty print all beamline states using rich.
         """
 
-        def _format_parameters(state_config: bl_states.BeamlineStateConfig) -> str:
-            parameter_dict = state_config.model_dump(exclude={"name"}, exclude_none=True)
-            if not parameter_dict:
-                return "-"
-            return "\n".join(f"{key}={value}" for key, value in parameter_dict.items())
-
         def _status_style(status_value: str) -> str:
             status_styles = {"valid": "green3", "invalid": "red3", "warning": "yellow3"}
             return status_styles.get(status_value.lower(), "grey50")
@@ -276,7 +342,8 @@ def _status_style(status_value: str) -> str:
         table.add_column("Label")
 
         for state in self._states.values():
-            params = _format_parameters(state)
+            state_class = _state_class_for_state_type(state.state_type)
+            params = state_class.format_config_summary(state)
             status = (
                 getattr(self, state.name).get()
                 if hasattr(self, state.name)
@@ -288,8 +355,8 @@ def _status_style(status_value: str) -> str:
                 str(state.name),
                 str(state.state_type),
                 str(params),
-                f"[{status_style}]{status_value}[/{status_style}]",
-                f"[{status_style}]{str(status.get('label', ''))}[/{status_style}]",
+                Text(status_value, style=status_style),
+                Text(_summarize_label(str(status.get("label", ""))), style=status_style),
             )
 
         console.print(table)