Skip to content

Commit f182538

Browse files
cappel89cappel89
committed
refactor: Improve config structure for aggregated state configs
1 parent 2a377f5 commit f182538

2 files changed

Lines changed: 109 additions & 27 deletions

File tree

‎bec_lib/bec_lib/bl_states.py‎

Lines changed: 95 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
import yaml
1313
from pydantic import BaseModel, field_validator, model_validator
14+
from typeguard import typechecked
1415

1516
from bec_lib import messages
1617
from bec_lib.alarm_handler import Alarms
@@ -132,14 +133,47 @@ class SignalConfig(BaseModel):
132133
abs_tol: float = 0.0
133134

134135

136+
class DeviceConfig(BaseModel):
137+
"""Configuration for a device inside a named machine state."""
138+
139+
abs_tol: float = 0.0
140+
value: float | int | str | bool | None = None
141+
low_limit: SignalConfig | None = None
142+
high_limit: SignalConfig | None = None
143+
signals: dict[str, SignalConfig] | None = None
144+
145+
@model_validator(mode="after")
146+
def validate_config(self) -> DeviceConfig:
147+
"""
148+
Validate that either value, low_limit, high_limit, or signals are provided.
149+
"""
150+
if (
151+
self.value is None
152+
and self.low_limit is None
153+
and self.high_limit is None
154+
and self.signals is None
155+
):
156+
raise ValueError(
157+
"At least one of value, low_limit, high_limit, or signals must be provided."
158+
)
159+
return self
160+
161+
135162
class SubDeviceStateConfig(BaseModel):
163+
"""
164+
Configuration for a sub-state with a specific label.
165+
This is a device/signal mappping to either a DeviceConfig or SignalConfig.
166+
"""
136167

137-
devices: dict[str, dict[str, SignalConfig]]
168+
devices: dict[str, DeviceConfig | SignalConfig]
169+
transition_metadata: dict[str, Any] | None = None
138170

139171

140172
class AggregatedStateConfig(BeamlineStateConfig):
141173
"""
142174
Configuration for a state machine driven by multiple device signals.
175+
176+
Keys of the states dictionary are the labels of the different states.
143177
"""
144178

145179
state_type: ClassVar[str] = "AggregatedState"
@@ -411,6 +445,7 @@ def _get_signal_source(self, signal_info: dict[str, Any]) -> SignalSource:
411445
f"{self._error_prefix} Unsupported kind: '{kind_str}' for signal : \n {yaml.dump(signal_info, indent=4)}"
412446
)
413447

448+
@typechecked
414449
def _resolve_signal(self, device_name: str, signal_name: str) -> tuple[str, SignalSource]:
415450
devices = self._get_devices()
416451
try:
@@ -466,25 +501,70 @@ def _build_rules(self) -> None:
466501
self._subscriptions.clear()
467502
for label, device_configs in self.config.states.items():
468503
state_requirements: list[_ResolvedStateSignal] = []
469-
for device_name, signal_configs in device_configs.devices.items():
470-
for signal_name, target in signal_configs.items():
471-
resolved_signal_name, source = self._resolve_signal(device_name, signal_name)
504+
for device_name, config in device_configs.devices.items():
505+
if isinstance(config, SignalConfig):
472506
state_requirements.append(
473-
_ResolvedStateSignal(
474-
label=label,
475-
device_name=device_name,
476-
signal_name=resolved_signal_name,
477-
expected_value=target.value,
478-
abs_tolerance=target.abs_tol,
479-
source=source,
507+
self._build_requirement_for_signal(
508+
device_name, device_name, config.value, config.abs_tol, label
480509
)
481510
)
482-
self._subscriptions.add((device_name, source))
483-
self._signal_info_to_labels.setdefault(
484-
(device_name, source, resolved_signal_name), set()
485-
).add(label)
511+
elif isinstance(config, DeviceConfig):
512+
# If a value is specified for the device, add it as a requirement
513+
if config.value is not None:
514+
state_requirements.append(
515+
self._build_requirement_for_signal(
516+
device_name, device_name, config.value, config.abs_tol, label
517+
)
518+
)
519+
if config.low_limit is not None:
520+
state_requirements.append(
521+
self._build_requirement_for_signal(
522+
device_name,
523+
"low_limit",
524+
config.low_limit.value,
525+
config.low_limit.abs_tol,
526+
label,
527+
)
528+
)
529+
if config.high_limit is not None:
530+
state_requirements.append(
531+
self._build_requirement_for_signal(
532+
device_name,
533+
"high_limit",
534+
config.high_limit.value,
535+
config.high_limit.abs_tol,
536+
label,
537+
)
538+
)
539+
for signal_name, signal_config in (config.signals or {}).items():
540+
state_requirements.append(
541+
self._build_requirement_for_signal(
542+
device_name,
543+
signal_name,
544+
signal_config.value,
545+
signal_config.abs_tol,
546+
label,
547+
)
548+
)
486549
self._requirements_for_label[label] = state_requirements
487550

551+
def _build_requirement_for_signal(
552+
self, device_name: str, signal_name: str, value: Any, abs_tol: float, label: str
553+
) -> _ResolvedStateSignal:
554+
resolved_signal_name, source = self._resolve_signal(device_name, signal_name)
555+
self._subscriptions.add((device_name, source))
556+
self._signal_info_to_labels.setdefault(
557+
(device_name, source, resolved_signal_name), set()
558+
).add(label)
559+
return _ResolvedStateSignal(
560+
label=label,
561+
device_name=device_name,
562+
signal_name=resolved_signal_name,
563+
expected_value=value,
564+
abs_tolerance=abs_tol,
565+
source=source,
566+
)
567+
488568
def start(self) -> None:
489569
if self.started:
490570
return

‎bec_lib/tests/test_beamline_states.py‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -208,27 +208,28 @@ def aggregated_state_config(self):
208208
"alignment": {
209209
"devices": {
210210
"samx": {
211-
"readback": {"value": 0, "abs_tol": 0.1},
212-
"velocity": {"value": 5, "abs_tol": 0.1},
211+
"value": 0,
212+
"abs_tol": 0.1,
213213
"low_limit": {"value": -20, "abs_tol": 0.1},
214214
"high_limit": {"value": 20, "abs_tol": 0.1},
215215
},
216-
"bpm4i": {"bpm4i": {"value": 0, "abs_tol": 0.1}},
216+
"bpm4i": {"value": 0, "abs_tol": 0.1},
217217
}
218218
},
219219
"measurement": {
220220
"devices": {
221221
"samx": {
222-
"readback": {"value": 19, "abs_tol": 0.1},
223-
"velocity": {"value": 5, "abs_tol": 0.1},
224-
"low_limit_travel": {"value": -20, "abs_tol": 0.1},
225-
"high_limit_travel": {"value": 20, "abs_tol": 0.1},
222+
"value": 19,
223+
"abs_tol": 0.1,
224+
"low_limit": {"value": -20, "abs_tol": 0.1},
225+
"high_limit": {"value": 20, "abs_tol": 0.1},
226+
"signals": {"velocity": {"value": 5, "abs_tol": 0.1}},
226227
},
227-
"bpm4i": {"bpm4i": {"value": 2, "abs_tol": 0.1}},
228+
"bpm4i": {"value": 2, "abs_tol": 0.1},
228229
}
229230
},
230-
"test": {"devices": {"bpm4i": {"bpm4i": {"value": 0, "abs_tol": 0.1}}}},
231-
"string_state": {"devices": {"bpm3i": {"bpm3i": {"value": "ok"}}}},
231+
"test": {"devices": {"bpm4i": {"value": 0, "abs_tol": 0.1}}},
232+
"string_state": {"devices": {"bpm3i": {"value": "ok"}}},
232233
},
233234
)
234235

@@ -647,8 +648,9 @@ def config_dict(self):
647648
"alignment": {
648649
"devices": {
649650
"samx": {
650-
"readback": {"value": 0, "abs_tol": 0.1},
651-
"velocity": {"value": 5, "abs_tol": 0.1},
651+
"value": 0,
652+
"abs_tol": 0.1,
653+
"signals": {"velocity": {"value": 5, "abs_tol": 0.1}},
652654
}
653655
}
654656
}

0 commit comments

Comments
 (0)