fix(typehints): fix typehints for panda_box and bec_signal

This commit is contained in:
2026-02-16 21:30:30 +01:00
parent ac6d781690
commit 016e830852
2 changed files with 34 additions and 29 deletions

View File

@@ -27,7 +27,7 @@ import os
import threading import threading
import uuid import uuid
from enum import StrEnum from enum import StrEnum
from typing import TYPE_CHECKING, Any, Callable, TypeAlias, Union from typing import TYPE_CHECKING, Any, Callable, TypeAlias, TypedDict
import pandablocks.commands as pbc import pandablocks.commands as pbc
from bec_lib import bec_logger from bec_lib import bec_logger
@@ -138,16 +138,21 @@ class PandaState(StrEnum):
# pylint: disable=invalid-name # pylint: disable=invalid-name
LITERAL_PANDA_COMMANDS: TypeAlias = Union[ LITERAL_PANDA_COMMANDS: TypeAlias = (
pbc.Raw, pbc.Raw
pbc.Arm, | pbc.Arm
pbc.Disarm, | pbc.Disarm
pbc.GetChanges, | pbc.GetChanges
pbc.GetBlockInfo, | pbc.GetBlockInfo
pbc.GetFieldInfo, | pbc.GetFieldInfo
pbc.GetPcapBitsLabels, | pbc.GetPcapBitsLabels
] )
LITERAL_PANDA_DATA: TypeAlias = Union[ReadyData, StartData, FrameData, EndData, Data] LITERAL_PANDA_DATA: TypeAlias = ReadyData | StartData | FrameData | EndData | Data
class DataCallback(TypedDict):
callback: Callable[[LITERAL_PANDA_DATA], None]
data_type: PandaState
class PandaBox(PSIDeviceBase): class PandaBox(PSIDeviceBase):
@@ -197,10 +202,10 @@ class PandaBox(PSIDeviceBase):
self._panda_state: PandaState | str = PandaState.DISARMED.value self._panda_state: PandaState | str = PandaState.DISARMED.value
# Status callback management # Status callback management
self._status_callbacks: dict[uuid.UUID, dict[str, Any]] = {} self._status_callbacks: dict[str, dict[str, Any]] = {}
# Data callbacks management # Data callbacks management
self._data_callbacks: dict[uuid.UUID, Callable[[LITERAL_PANDA_DATA], None]] = {} self._data_callbacks: dict[str, DataCallback] = {}
# Thread to receive data from the PandaBox # Thread to receive data from the PandaBox
self.data_thread: threading.Thread = threading.Thread( self.data_thread: threading.Thread = threading.Thread(
@@ -226,7 +231,7 @@ class PandaBox(PSIDeviceBase):
### Public API methods ### ### Public API methods ###
########################## ##########################
def send_raw(self, raw_command: Union[str, list[str]]) -> Any: def send_raw(self, raw_command: str | list[str]) -> Any:
""" """
Send a raw command to the PandaBox. This can be used to set for example Send a raw command to the PandaBox. This can be used to set for example
values on PandaBox block fields directly, e.g. 'BITS.B=1' to set the BITS.B field to 1. values on PandaBox block fields directly, e.g. 'BITS.B=1' to set the BITS.B field to 1.
@@ -245,6 +250,8 @@ class PandaBox(PSIDeviceBase):
- ['PULSE1.DELAY.UNITS=s', PULSE1.DELAY=0, PULSE1.WIDTH.UNITS=s, PULSE1.WIDTH=0.001] to set multiple fields at once - ['PULSE1.DELAY.UNITS=s', PULSE1.DELAY=0, PULSE1.WIDTH.UNITS=s, PULSE1.WIDTH=0.001] to set multiple fields at once
- '*CAPTURE?' to inspect which signals have been configured for capture (PCAP?) TODO to check - '*CAPTURE?' to inspect which signals have been configured for capture (PCAP?) TODO to check
""" """
if isinstance(raw_command, str):
raw_command = [raw_command]
return self._send_command(pbc.Raw(raw_command)) return self._send_command(pbc.Raw(raw_command))
def add_status_callback( def add_status_callback(
@@ -310,7 +317,7 @@ class PandaBox(PSIDeviceBase):
def add_data_callback( def add_data_callback(
self, self,
callback: Callable[[LITERAL_PANDA_DATA], None], callback: Callable[[LITERAL_PANDA_DATA], None],
data_type: PandaState = PandaState.FRAME.value, data_type: PandaState = PandaState.FRAME,
) -> str: ) -> str:
""" """
Register a data callback to be called whenever new data is received from the PandaBox. Register a data callback to be called whenever new data is received from the PandaBox.
@@ -406,19 +413,19 @@ class PandaBox(PSIDeviceBase):
for data in client.data(scaled=False): for data in client.data(scaled=False):
if isinstance(data, ReadyData): if isinstance(data, ReadyData):
self._run_status_callbacks(PandaState.READY) self._run_status_callbacks(PandaState.READY)
self._run_data_callbacks(data, PandaState.READY.value) self._run_data_callbacks(data, PandaState.READY)
elif isinstance(data, StartData): elif isinstance(data, StartData):
self._run_status_callbacks(PandaState.START) self._run_status_callbacks(PandaState.START)
self._run_data_callbacks(data, PandaState.START.value) self._run_data_callbacks(data, PandaState.START)
elif isinstance(data, FrameData): elif isinstance(data, FrameData):
self._run_status_callbacks(PandaState.FRAME) self._run_status_callbacks(PandaState.FRAME)
self._run_data_callbacks(data, PandaState.FRAME.value) self._run_data_callbacks(data, PandaState.FRAME)
elif isinstance(data, EndData): elif isinstance(data, EndData):
self._run_status_callbacks(PandaState.END) self._run_status_callbacks(PandaState.END)
self._run_data_callbacks(data, PandaState.END.value) self._run_data_callbacks(data, PandaState.END)
break # Exit data readout loop break # Exit data readout loop
finally: finally:
@@ -441,7 +448,7 @@ class PandaBox(PSIDeviceBase):
# As DISARMED is not triggered by a data message, we manually run data callbacks for it here # As DISARMED is not triggered by a data message, we manually run data callbacks for it here
# and run it with an empty Data() object following the base class for data message responses # and run it with an empty Data() object following the base class for data message responses
# of the pandablocks library. # of the pandablocks library.
self._run_data_callbacks(Data(), PandaState.DISARMED.value) self._run_data_callbacks(Data(), PandaState.DISARMED)
def _run_status_callbacks(self, event: PandaState) -> None: def _run_status_callbacks(self, event: PandaState) -> None:
""" """
@@ -522,7 +529,7 @@ class PandaBox(PSIDeviceBase):
raise e from e raise e from e
super().on_connected() super().on_connected()
self.data_thread.start() self.data_thread.start()
self.add_data_callback(data_type=PandaState.FRAME.value, callback=self._receive_frame_data) self.add_data_callback(data_type=PandaState.FRAME, callback=self._receive_frame_data)
def _receive_frame_data(self, data: FrameData) -> None: def _receive_frame_data(self, data: FrameData) -> None:
logger.info(f"Received frame data with signals {data}") logger.info(f"Received frame data with signals {data}")

View File

@@ -79,6 +79,9 @@ class SignalInfo(BaseModel):
) )
_SignalsTypes = list[tuple[str, str | Kind]] | list[str] | str | None
class BECMessageSignal(Signal): class BECMessageSignal(Signal):
""" """
Custom signal class that accepts BECMessage objects as values. Custom signal class that accepts BECMessage objects as values.
@@ -98,12 +101,7 @@ class BECMessageSignal(Signal):
role: Literal["main", "preview", "diagnostic", "file event", "progress"] = "main", role: Literal["main", "preview", "diagnostic", "file event", "progress"] = "main",
acquisition_group: Literal["baseline", "monitored"] | str | None = None, acquisition_group: Literal["baseline", "monitored"] | str | None = None,
enabled: bool = True, enabled: bool = True,
signals: ( signals: _SignalsTypes | Callable[[], _SignalsTypes] = None,
Callable[[], list[str]]
| Callable[[], list[tuple[str, str | Kind]]]
| list[tuple[str, str | Kind] | str]
| None
) = None,
signal_metadata: dict | None = None, signal_metadata: dict | None = None,
**kwargs, **kwargs,
): ):
@@ -140,7 +138,7 @@ class BECMessageSignal(Signal):
self._bec_message_type = bec_message_type self._bec_message_type = bec_message_type
def _unify_signals( def _unify_signals(
self, signals: Callable[[], list[str]] | list[tuple[str, str | Kind] | str] | str | None self, signals: _SignalsTypes | Callable[[], _SignalsTypes]
) -> list[tuple[str, int]]: ) -> list[tuple[str, int]]:
""" """
Unify the signals list to a list of tuples with signal name and kind. Unify the signals list to a list of tuples with signal name and kind.
@@ -151,7 +149,7 @@ class BECMessageSignal(Signal):
Returns: Returns:
list[tuple[str, str]]: The unified list of signals. list[tuple[str, str]]: The unified list of signals.
""" """
if isinstance(signals, Callable): if callable(signals):
signals = signals() signals = signals()
if signals is None: if signals is None:
return [(self.name, Kind.hinted.value)] # Default to signal name with hinted kind return [(self.name, Kind.hinted.value)] # Default to signal name with hinted kind