Source code for ophyd_async.core.signal

from __future__ import annotations

import asyncio
import functools
from typing import (

from bluesky.protocols import (

from ophyd_async.core.mock_signal_backend import MockSignalBackend
from ophyd_async.protocols import AsyncConfigurable, AsyncReadable, AsyncStageable

from .async_status import AsyncStatus
from .device import Device
from .signal_backend import SignalBackend
from .soft_signal_backend import SignalMetadata, SoftSignalBackend
from .utils import DEFAULT_TIMEOUT, CalculatableTimeout, CalculateTimeout, Callback, T

def _add_timeout(func):
    async def wrapper(self: Signal, *args, **kwargs):
        return await asyncio.wait_for(func(self, *args, **kwargs), self._timeout)

    return wrapper

def _fail(self, other, *args, **kwargs):
    if isinstance(other, Signal):
        raise TypeError(
            "Can't compare two Signals, did you mean await signal.get_value() instead?"
        return NotImplemented

[docs] class Signal(Device, Generic[T]): """A Device with the concept of a value, with R, RW, W and X flavours""" def __init__( self, backend: Optional[SignalBackend[T]] = None, timeout: Optional[float] = DEFAULT_TIMEOUT, name: str = "", ) -> None: self._timeout = timeout self._backend = backend super().__init__(name)
[docs] async def connect( self, mock=False, timeout=DEFAULT_TIMEOUT, force_reconnect: bool = False, backend: Optional[SignalBackend[T]] = None, ): if backend: if self._backend and backend is not self._backend: raise ValueError("Backend at connection different from previous one.") self._backend = backend if ( self._previous_connect_was_mock is not None and self._previous_connect_was_mock != mock ): raise RuntimeError( f"`connect(mock={mock})` called on a `Signal` where the previous " f"connect was `mock={self._previous_connect_was_mock}`. Changing mock " "value between connects is not permitted." ) self._previous_connect_was_mock = mock if mock and not issubclass(type(self._backend), MockSignalBackend): # Using a soft backend, look to the initial value self._backend = MockSignalBackend(initial_backend=self._backend) if self._backend is None: raise RuntimeError("`connect` called on signal without backend") can_use_previous_connection: bool = self._connect_task is not None and not ( self._connect_task.done() and self._connect_task.exception() ) if force_reconnect or not can_use_previous_connection: self.log.debug(f"Connecting to {self.source}") self._connect_task = asyncio.create_task( self._backend.connect(timeout=timeout) ) else: self.log.debug(f"Reusing previous connection to {self.source}") assert ( self._connect_task ), "this assert is for type analysis and will never fail" await self._connect_task
@property def source(self) -> str: """Like ca://PV_PREFIX:SIGNAL, or "" if not set""" return self._backend.source( __lt__ = __le__ = __eq__ = __ge__ = __gt__ = __ne__ = _fail def __hash__(self): # Restore the default implementation so we can use in a set or dict return hash(id(self))
class _SignalCache(Generic[T]): def __init__(self, backend: SignalBackend[T], signal: Signal): self._signal = signal self._staged = False self._listeners: Dict[Callback, bool] = {} self._valid = asyncio.Event() self._reading: Optional[Reading] = None self._value: Optional[T] = None self.backend = backend signal.log.debug(f"Making subscription on source {signal.source}") backend.set_callback(self._callback) def close(self): self.backend.set_callback(None) self._signal.log.debug(f"Closing subscription on source {self._signal.source}") async def get_reading(self) -> Reading: await self._valid.wait() assert self._reading is not None, "Monitor not working" return self._reading async def get_value(self) -> T: await self._valid.wait() assert self._value is not None, "Monitor not working" return self._value def _callback(self, reading: Reading, value: T): self._signal.log.debug( f"Updated subscription: reading of source {self._signal.source} changed" f"from {self._reading} to {reading}" ) self._reading = reading self._value = value self._valid.set() for function, want_value in self._listeners.items(): self._notify(function, want_value) def _notify(self, function: Callback, want_value: bool): if want_value: function(self._value) else: function({ self._reading}) def subscribe(self, function: Callback, want_value: bool) -> None: self._listeners[function] = want_value if self._valid.is_set(): self._notify(function, want_value) def unsubscribe(self, function: Callback) -> bool: self._listeners.pop(function) return self._staged or bool(self._listeners) def set_staged(self, staged: bool): self._staged = staged return self._staged or bool(self._listeners)
[docs] class SignalR(Signal[T], AsyncReadable, AsyncStageable, Subscribable): """Signal that can be read from and monitored""" _cache: Optional[_SignalCache] = None def _backend_or_cache( self, cached: Optional[bool] ) -> Union[_SignalCache, SignalBackend]: # If cached is None then calculate it based on whether we already have a cache if cached is None: cached = self._cache is not None if cached: assert self._cache, f"{self.source} not being monitored" return self._cache else: return self._backend def _get_cache(self) -> _SignalCache: if not self._cache: self._cache = _SignalCache(self._backend, self) return self._cache def _del_cache(self, needed: bool): if self._cache and not needed: self._cache.close() self._cache = None
[docs] @_add_timeout async def read(self, cached: Optional[bool] = None) -> Dict[str, Reading]: """Return a single item dict with the reading in it""" return { await self._backend_or_cache(cached).get_reading()}
[docs] @_add_timeout async def describe(self) -> Dict[str, DataKey]: """Return a single item dict with the descriptor in it""" return { await self._backend.get_datakey(self.source)}
[docs] @_add_timeout async def get_value(self, cached: Optional[bool] = None) -> T: """The current value""" value = await self._backend_or_cache(cached).get_value() self.log.debug(f"get_value() on source {self.source} returned {value}") return value
[docs] def subscribe_value(self, function: Callback[T]): """Subscribe to updates in value of a device""" self._get_cache().subscribe(function, want_value=True)
[docs] def subscribe(self, function: Callback[Dict[str, Reading]]) -> None: """Subscribe to updates in the reading""" self._get_cache().subscribe(function, want_value=False)
[docs] def clear_sub(self, function: Callback) -> None: """Remove a subscription.""" self._del_cache(self._get_cache().unsubscribe(function))
[docs] @AsyncStatus.wrap async def stage(self) -> None: """Start caching this signal""" self._get_cache().set_staged(True)
[docs] @AsyncStatus.wrap async def unstage(self) -> None: """Stop caching this signal""" self._del_cache(self._get_cache().set_staged(False))
[docs] class SignalW(Signal[T], Movable): """Signal that can be set"""
[docs] def set( self, value: T, wait=True, timeout: CalculatableTimeout = CalculateTimeout ) -> AsyncStatus: """Set the value and return a status saying when it's done""" if timeout is CalculateTimeout: timeout = self._timeout async def do_set(): self.log.debug(f"Putting value {value} to backend at source {self.source}") await self._backend.put(value, wait=wait, timeout=timeout) self.log.debug( f"Successfully put value {value} to backend at source {self.source}" ) return AsyncStatus(do_set())
[docs] class SignalRW(SignalR[T], SignalW[T], Locatable): """Signal that can be both read and set"""
[docs] async def locate(self) -> Location: location: Location = { "setpoint": await self._backend.get_setpoint(), "readback": await self.get_value(), } return location
[docs] class SignalX(Signal): """Signal that puts the default value"""
[docs] def trigger( self, wait=True, timeout: CalculatableTimeout = CalculateTimeout ) -> AsyncStatus: """Trigger the action and return a status saying when it's done""" if timeout is CalculateTimeout: timeout = self._timeout coro = self._backend.put(None, wait=wait, timeout=timeout) return AsyncStatus(coro)
[docs] def soft_signal_rw( datatype: Optional[Type[T]] = None, initial_value: Optional[T] = None, name: str = "", units: str | None = None, precision: int | None = None, ) -> SignalRW[T]: """Creates a read-writable Signal with a SoftSignalBackend. May pass metadata, which are propagated into describe. """ metadata = SignalMetadata(units=units, precision=precision) signal = SignalRW( SoftSignalBackend(datatype, initial_value, metadata=metadata), name=name, ) return signal
[docs] def soft_signal_r_and_setter( datatype: Optional[Type[T]] = None, initial_value: Optional[T] = None, name: str = "", units: str | None = None, precision: int | None = None, ) -> Tuple[SignalR[T], Callable[[T], None]]: """Returns a tuple of a read-only Signal and a callable through which the signal can be internally modified within the device. May pass metadata, which are propagated into describe. Use soft_signal_rw if you want a device that is externally modifiable """ metadata = SignalMetadata(units=units, precision=precision) backend = SoftSignalBackend(datatype, initial_value, metadata=metadata) signal = SignalR(backend, name=name) return (signal, backend.set_value)
def _generate_assert_error_msg( name: str, expected_result: str, actual_result: str ) -> str: WARNING = "\033[93m" FAIL = "\033[91m" ENDC = "\033[0m" return ( f"Expected {WARNING}{name}{ENDC} to produce" + f"\n{FAIL}{expected_result}{ENDC}" + f"\nbut actually got \n{FAIL}{actual_result}{ENDC}" )
[docs] async def assert_value(signal: SignalR[T], value: Any) -> None: """Assert a signal's value and compare it an expected signal. Parameters ---------- signal: signal with get_value. value: The expected value from the signal. Notes ----- Example usage:: await assert_value(signal, value) """ actual_value = await signal.get_value() assert actual_value == value, _generate_assert_error_msg(, expected_result=value, actual_result=actual_value, )
[docs] async def assert_reading( readable: AsyncReadable, expected_reading: Mapping[str, Reading] ) -> None: """Assert readings from readable. Parameters ---------- readable: Callable with function that generate readings. reading: The expected readings from the readable. Notes ----- Example usage:: await assert_reading(readable, reading) """ actual_reading = await assert expected_reading == actual_reading, _generate_assert_error_msg(, expected_result=expected_reading, actual_result=actual_reading, )
[docs] async def assert_configuration( configurable: AsyncConfigurable, configuration: Mapping[str, Reading], ) -> None: """Assert readings from Configurable. Parameters ---------- configurable: Configurable with function that generate readings. configuration: The expected readings from configurable. Notes ----- Example usage:: await assert_configuration(configurable configuration) """ actual_configurable = await configurable.read_configuration() assert configuration == actual_configurable, _generate_assert_error_msg(, expected_result=configuration, actual_result=actual_configurable, )
[docs] def assert_emitted(docs: Mapping[str, list[dict]], **numbers: int): """Assert emitted document generated by running a Bluesky plan Parameters ---------- Doc: A dictionary numbers: expected emission in kwarg from Notes ----- Example usage:: assert_emitted(docs, start=1, descriptor=1, resource=1, datum=1, event=1, stop=1) """ assert list(docs) == list(numbers), _generate_assert_error_msg( name="documents", expected_result=list(numbers), actual_result=list(docs), ) actual_numbers = {name: len(d) for name, d in docs.items()} assert actual_numbers == numbers, _generate_assert_error_msg( name="emitted", expected_result=numbers, actual_result=actual_numbers, )
[docs] async def observe_value( signal: SignalR[T], timeout: float | None = None, done_status: Status | None = None ) -> AsyncGenerator[T, None]: """Subscribe to the value of a signal so it can be iterated from. Parameters ---------- signal: Call subscribe_value on this at the start, and clear_sub on it at the end timeout: If given, how long to wait for each updated value in seconds. If an update is not produced in this time then raise asyncio.TimeoutError done_status: If this status is complete, stop observing and make the iterator return. If it raises an exception then this exception will be raised by the iterator. Notes ----- Example usage:: async for value in observe_value(sig): do_something_with(value) """ q: asyncio.Queue[T | Status] = asyncio.Queue() if timeout is None: get_value = q.get else: async def get_value(): return await asyncio.wait_for(q.get(), timeout) if done_status is not None: done_status.add_callback(q.put_nowait) signal.subscribe_value(q.put_nowait) try: while True: item = await get_value() if done_status and item is done_status: if exc := done_status.exception(): raise exc else: break else: yield item finally: signal.clear_sub(q.put_nowait)
class _ValueChecker(Generic[T]): def __init__(self, matcher: Callable[[T], bool], matcher_name: str): self._last_value: Optional[T] = None self._matcher = matcher self._matcher_name = matcher_name async def _wait_for_value(self, signal: SignalR[T]): async for value in observe_value(signal): self._last_value = value if self._matcher(value): return async def wait_for_value(self, signal: SignalR[T], timeout: Optional[float]): try: await asyncio.wait_for(self._wait_for_value(signal), timeout) except asyncio.TimeoutError as e: raise TimeoutError( f"{} didn't match {self._matcher_name} in {timeout}s, " f"last value {self._last_value!r}" ) from e
[docs] async def wait_for_value( signal: SignalR[T], match: Union[T, Callable[[T], bool]], timeout: Optional[float] ): """Wait for a signal to have a matching value. Parameters ---------- signal: Call subscribe_value on this at the start, and clear_sub on it at the end match: If a callable, it should return True if the value matches. If not callable then value will be checked for equality with match. timeout: How long to wait for the value to match Notes ----- Example usage:: wait_for_value(device.acquiring, 1, timeout=1) Or:: wait_for_value(device.num_captured, lambda v: v > 45, timeout=1) """ if callable(match): checker = _ValueChecker(match, match.__name__) else: checker = _ValueChecker(lambda v: v == match, repr(match)) await checker.wait_for_value(signal, timeout)
[docs] async def set_and_wait_for_value( signal: SignalRW[T], value: T, timeout: float = DEFAULT_TIMEOUT, status_timeout: Optional[float] = None, ) -> AsyncStatus: """Set a signal and monitor it until it has that value. Useful for busy record, or other Signals with pattern: - Set Signal with wait=True and stash the Status - Read the same Signal to check the operation has started - Return the Status so calling code can wait for operation to complete This function sets a signal to a specified value, optionally with or without a ca/pv put callback, and waits for the readback value of the signal to match the value it was set to. Parameters ---------- signal: The signal to set and monitor value: The value to set it to timeout: How long to wait for the signal to have the value status_timeout: How long the returned Status will wait for the set to complete Notes ----- Example usage:: set_and_wait_for_value(device.acquire, 1) """ status = signal.set(value, timeout=status_timeout) await wait_for_value(signal, value, timeout=timeout) return status