fix(beamline-states): better pydantic model handling

Author: wyzula-jan

bec_widgets/widgets/services/beamline_states/beamline_state_manager.py

@@ -3,10 +3,9 @@
 import sys
 from typing import Any
 
-from bec_lib import bl_states
+from bec_lib import bl_states, messages
 from bec_lib.endpoints import MessageEndpoints
 from bec_qthemes import material_icon
-from pydantic import BaseModel
 from qtpy.QtCore import QAbstractListModel, QModelIndex, QSize, Qt
 from qtpy.QtWidgets import (
     QAbstractItemView,
@@ -48,7 +47,7 @@ def __init__(self, parent: QWidget | None = None) -> None:
         super().__init__(parent)
         self._state_order: list[str] = []
         self._state_rows: dict[str, int] = {}
-        self._state_configs: dict[str, dict[str, Any]] = {}
+        self._state_configs: dict[str, messages.BeamlineStateConfig] = {}
 
     def rowCount(self, parent: QModelIndex = QModelIndex()) -> int:  # noqa: N802
         return 0 if parent.isValid() else len(self._state_order)
@@ -60,17 +59,17 @@ def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> A
         if role in (Qt.ItemDataRole.DisplayRole, self.NameRole):
             return name
         if role == self.ConfigRole:
-            return self._state_configs.get(name, {})
+            return self._state_configs.get(name)
         return None
 
     def flags(self, index: QModelIndex) -> Qt.ItemFlag:
         if not index.isValid():
             return Qt.ItemFlag.NoItemFlags
         return Qt.ItemFlag.ItemIsEnabled | Qt.ItemFlag.ItemIsSelectable
 
-    def set_states(self, state_configs: list[dict[str, Any]]) -> None:
-        new_order = [str(state["name"]) for state in state_configs if state.get("name")]
-        new_configs = {str(state["name"]): state for state in state_configs if state.get("name")}
+    def set_states(self, state_configs: list[messages.BeamlineStateConfig]) -> None:
+        new_order = [state.name for state in state_configs]
+        new_configs = {state.name: state for state in state_configs}
 
         for row in reversed(
             [row for row, name in enumerate(self._state_order) if name not in new_configs]
@@ -129,11 +128,8 @@ def createEditor(  # noqa: N802
         self, parent: QWidget, _option: QStyleOptionViewItem, index: QModelIndex
     ) -> QWidget:
         name = index.data(_BeamlineStateListModel.NameRole)
-        state_config = index.data(_BeamlineStateListModel.ConfigRole) or {}
-        title = state_config.get("title") or name
-        pill = BeamlineStatePill(
-            parent=parent, state_name=name, title=title, client=self._manager.client
-        )
+        state_config = index.data(_BeamlineStateListModel.ConfigRole)
+        pill = BeamlineStatePill(parent=parent, state_name=name, client=self._manager.client)
         pill.idle_card_background = self._manager.idle_card_background
         pill.set_state_config(state_config)
         pill.state_changed.connect(self._manager._on_pill_state_changed)
@@ -147,9 +143,8 @@ def setEditorData(self, editor: QWidget, index: QModelIndex) -> None:  # noqa: N
         if not isinstance(editor, BeamlineStatePill):
             return
         name = index.data(_BeamlineStateListModel.NameRole)
-        state_config = index.data(_BeamlineStateListModel.ConfigRole) or {}
-        title = state_config.get("title") or name
-        editor.set_state_name(str(name), title=str(title))
+        state_config = index.data(_BeamlineStateListModel.ConfigRole)
+        editor.set_state_name(str(name))
         editor.idle_card_background = self._manager.idle_card_background
         editor.set_state_config(state_config)
 
@@ -229,7 +224,7 @@ def __init__(
         )
         self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
         self._state_pills: dict[str, BeamlineStatePill] = {}
-        self._state_configs: dict[str, dict[str, Any]] = {}
+        self._state_configs: dict[str, messages.BeamlineStateConfig] = {}
         self._state_order: list[str] = []
         self._selected_statuses: set[str] | None = None
         self._selected_devices: set[str] | None = None
@@ -414,14 +409,12 @@ def update_available_states(
     ) -> None:
         """Update the displayed pills from ``AvailableBeamlineStatesMessage`` content."""
         expanded_names = {name for name, pill in self._state_pills.items() if pill.is_expanded()}
-        states = content.get("states", [])
-        state_configs = [self._state_config_to_dict(state) for state in states]
-        state_configs = [state for state in state_configs if state.get("name")]
+        state_configs: list[messages.BeamlineStateConfig] = content.get("states", [])
         if state_configs == list(self._state_configs.values()):
             self._apply_filters()
             return
-        self._state_configs = {str(state["name"]): state for state in state_configs}
-        self._state_order = [str(state["name"]) for state in state_configs]
+        self._state_configs = {state.name: state for state in state_configs}
+        self._state_order = [state.name for state in state_configs]
         self._model.set_states(state_configs)
         self._open_persistent_editors(expanded_names)
         self._apply_filters()
@@ -472,7 +465,7 @@ def _is_state_visible(self, name: str) -> bool:
         ):
             return False
 
-        device = self._state_device(self._state_configs.get(name, {}))
+        device = self._state_device(self._state_configs.get(name))
         if self._selected_devices is not None and device not in self._selected_devices:
             return False
 
@@ -505,14 +498,13 @@ def _update_state_parameters(
             )
             return
         try:
-            parameters = config.model_dump(exclude={"name"})
-            state_client.update_parameters(**parameters)
+            state_client.update_parameters(**config.model_dump(exclude={"name"}))
         except Exception as exc:
             QMessageBox.warning(self, "Cannot Update State", str(exc))
             return
         pill = self._state_pills.get(state_name)
         if pill is not None:
-            pill.mark_current_settings_clean(config)
+            pill.mark_current_settings_clean()
 
     @SafeSlot(str)
     def _remove_state_requested(self, state_name: str) -> None:
@@ -551,27 +543,10 @@ def _refresh_hidden_summary(self, hidden_count: int | None = None) -> None:
         )
 
     @staticmethod
-    def _state_device(state: dict[str, Any]) -> str | None:
-        parameters = state.get("parameters")
-        if isinstance(parameters, dict):
-            device = parameters.get("device")
-        else:
-            device = state.get("device")
+    def _state_device(state: messages.BeamlineStateConfig | None) -> str | None:
+        device = state.parameters.get("device") if state is not None else None
         return str(device) if device else None
 
-    @staticmethod
-    def _state_config_to_dict(state: dict[str, Any] | BaseModel) -> dict[str, Any]:
-        if isinstance(state, dict):
-            return state
-        state_dict = state.model_dump()
-        state_type = getattr(state, "state_type", None)
-        if state_type is not None:
-            state_dict.setdefault("state_type", state_type)
-        title = getattr(state, "title", None)
-        if title is not None and not state_dict.get("title"):
-            state_dict["title"] = title
-        return state_dict
-
     def cleanup(self) -> None:
         self.bec_dispatcher.disconnect_slot(
             self.update_available_states, MessageEndpoints.available_beamline_states()

bec_widgets/widgets/services/beamline_states/beamline_state_pill.py

@@ -2,7 +2,7 @@
 
 from typing import Any
 
-from bec_lib import bl_states
+from bec_lib import bl_states, messages
 from bec_lib.endpoints import MessageEndpoints
 from bec_qthemes import material_icon
 from qtpy.QtCore import Qt, Signal
@@ -76,7 +76,6 @@ def __init__(
         self,
         parent: QWidget | None = None,
         state_name: str | None = None,
-        title: str | None = None,
         client=None,
         config: ConnectionConfig | None = None,
         gui_id: str | None = None,
@@ -89,8 +88,7 @@ def __init__(
         self.setAttribute(Qt.WidgetAttribute.WA_StyledBackground, True)
         self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Minimum)
         self._state_name: str | None = None
-        self._title: str | None = None
-        self._state_config: dict[str, Any] = {}
+        self._state_config: messages.BeamlineStateConfig | None = None
         self._status = "unknown"
         self._label = "No state information available."
         self._expanded = False
@@ -100,9 +98,9 @@ def __init__(
         self._settings_dirty_fields: set[str] = set()
         self._settings_form_stale = True
 
-        self._init_ui(state_name, title)
+        self._init_ui(state_name)
 
-    def _init_ui(self, state_name: str | None = None, title: str | None = None) -> None:
+    def _init_ui(self, state_name: str | None = None) -> None:
         self._shadow = QGraphicsDropShadowEffect(self)
         self._shadow.setBlurRadius(18)
         self._shadow.setOffset(0, 2)
@@ -212,7 +210,7 @@ def _init_ui(self, state_name: str | None = None, title: str | None = None) -> N
         self.setLayout(layout)
 
         self._settings.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Minimum)
-        self.set_state_name(state_name, title=title)
+        self.set_state_name(state_name)
         self._update_button.setEnabled(False)
         self._revert_button.setEnabled(False)
 
@@ -225,15 +223,14 @@ def state_name(self) -> str | None:
     def state_name(self, state_name: str | None) -> None:
         self.set_state_name(state_name)
 
-    def set_state_name(self, state_name: str | None, title: str | None = None) -> None:
+    def set_state_name(self, state_name: str | None) -> None:
         """
         Set the BEC beamline state this pill displays.
 
         Args:
             state_name: State name as published by ``AvailableBeamlineStatesMessage``.
-            title: Optional human-readable title for the state.
         """
-        if state_name == self._state_name and title == self._title:
+        if state_name == self._state_name:
             return
 
         if self._state_name is not None:
@@ -242,8 +239,7 @@ def set_state_name(self, state_name: str | None, title: str | None = None) -> No
             )
 
         self._state_name = state_name
-        self._title = title
-        self._name_label.setText(title or state_name or "Beamline state")
+        self._name_label.setText(state_name or "Beamline state")
 
         if self._state_name is None:
             self._set_visual_state("unknown", "No beamline state selected.")
@@ -255,13 +251,13 @@ def set_state_name(self, state_name: str | None, title: str | None = None) -> No
             self.update_state, MessageEndpoints.beamline_state(self._state_name)
         )
 
-    def set_state_config(self, state_config: dict[str, Any]) -> None:
+    def set_state_config(self, state_config: messages.BeamlineStateConfig | None) -> None:
         """Set the editable BEC state configuration displayed by the expanded panel."""
         self._state_config = state_config
         self._settings_form_stale = True
         if self._config_form is not None:
             self._populate_settings()
-            self._mark_settings_clean_from_current()
+            self.mark_current_settings_clean()
 
     @SafeProperty(bool, default=False)
     def idle_card_background(self) -> bool:
@@ -439,13 +435,13 @@ def _ensure_config_form(
     def _ensure_settings_form_current(self) -> PydanticWidgetForm:
         if self._settings_form_stale:
             self._populate_settings()
-            self._mark_settings_clean_from_current()
+            self.mark_current_settings_clean()
         return self._ensure_config_form()
 
     def _populate_settings(self) -> None:
         self._populating_settings = True
         try:
-            state_type = str(self._state_config.get("state_type") or "")
+            state_type = self._state_config.state_type if self._state_config is not None else ""
             config_class = None
             for state_class in SUPPORTED_BEAMLINE_STATES:
                 if state_type in {state_class.__name__, state_class.CONFIG_CLASS.state_type}:
@@ -468,25 +464,8 @@ def edited_config(self) -> bl_states.BeamlineStateConfig:
         config = self._ensure_settings_form_current().model_instance()
         return config  # type: ignore[return-value]
 
-    def mark_current_settings_clean(
-        self, config: bl_states.BeamlineStateConfig | None = None
-    ) -> None:
+    def mark_current_settings_clean(self) -> None:
         """Mark the current editor values as saved."""
-        if config is None:
-            parameters = self._ensure_config_form().raw_editable_data()
-        else:
-            parameters = config.model_dump(exclude={"name"})
-        if self._state_config:
-            state_parameters = self._state_config.get("parameters")
-            if isinstance(state_parameters, dict):
-                state_parameters.update(parameters)
-            else:
-                self._state_config.update(parameters)
-            if "title" in parameters:
-                self._state_config["title"] = parameters["title"]
-        self._mark_settings_clean_from_current()
-
-    def _mark_settings_clean_from_current(self) -> None:
         config_form = self._ensure_config_form()
         self._settings_baseline = config_form.raw_editable_data()
         config_form.mark_clean()
@@ -561,15 +540,12 @@ def _state_data_for_form(
         self, config_class: type[bl_states.BeamlineStateConfig]
     ) -> dict[str, Any]:
         data: dict[str, Any] = {}
-        parameters = self._state_config.get("parameters")
-        parameter_values = parameters if isinstance(parameters, dict) else {}
+        parameters = self._state_config.parameters if self._state_config is not None else {}
         for name in config_class.model_fields:
             if name == "name":
                 data[name] = self._state_name
-            elif name in parameter_values:
-                data[name] = parameter_values[name]
-            elif name in self._state_config:
-                data[name] = self._state_config[name]
+            elif name in parameters:
+                data[name] = parameters[name]
         return data
 
     @staticmethod