Source code for ophyd_async.core.device

"""Base device"""

from __future__ import annotations

import sys
from typing import (
    Any,
    Coroutine,
    Dict,
    Generator,
    Iterator,
    Optional,
    Set,
    Tuple,
    TypeVar,
)

from bluesky.protocols import HasName
from bluesky.run_engine import call_in_bluesky_event_loop

from .utils import DEFAULT_TIMEOUT, NotConnected, wait_for_connection


[docs] class Device(HasName): """Common base class for all Ophyd Async Devices. By default, names and connects all Device children. """ _name: str = "" #: The parent Device if it exists parent: Optional[Device] = None def __init__(self, name: str = "") -> None: self.set_name(name) @property def name(self) -> str: """Return the name of the Device""" return self._name def children(self) -> Iterator[Tuple[str, Device]]: for attr_name, attr in self.__dict__.items(): if attr_name != "parent" and isinstance(attr, Device): yield attr_name, attr
[docs] def set_name(self, name: str): """Set ``self.name=name`` and each ``self.child.name=name+"-child"``. Parameters ---------- name: New name to set """ self._name = name for attr_name, child in self.children(): child_name = f"{name}-{attr_name.rstrip('_')}" if name else "" child.set_name(child_name) child.parent = self
[docs] async def connect(self, sim: bool = False, timeout: float = DEFAULT_TIMEOUT): """Connect self and all child Devices. Contains a timeout that gets propagated to child.connect methods. Parameters ---------- sim: If True then connect in simulation mode. timeout: Time to wait before failing with a TimeoutError. """ coros = { name: child_device.connect(sim, timeout=timeout) for name, child_device in self.children() } if coros: await wait_for_connection(**coros)
VT = TypeVar("VT", bound=Device)
[docs] class DeviceVector(Dict[int, VT], Device): def children(self) -> Generator[Tuple[str, Device], None, None]: for attr_name, attr in self.items(): if isinstance(attr, Device): yield str(attr_name), attr
[docs] class DeviceCollector: """Collector of top level Device instances to be used as a context manager Parameters ---------- set_name: If True, call ``device.set_name(variable_name)`` on all collected Devices connect: If True, call ``device.connect(sim)`` in parallel on all collected Devices sim: If True, connect Signals in simulation mode timeout: How long to wait for connect before logging an exception Notes ----- Example usage:: [async] with DeviceCollector(): t1x = motor.Motor("BLxxI-MO-TABLE-01:X") t1y = motor.Motor("pva://BLxxI-MO-TABLE-01:Y") # Names and connects devices here assert t1x.comm.velocity.source assert t1x.name == "t1x" """ def __init__( self, set_name=True, connect=True, sim=False, timeout: float = 10.0, ): self._set_name = set_name self._connect = connect self._sim = sim self._timeout = timeout self._names_on_enter: Set[str] = set() self._objects_on_exit: Dict[str, Any] = {} def _caller_locals(self): """Walk up until we find a stack frame that doesn't have us as self""" try: raise ValueError except ValueError: _, _, tb = sys.exc_info() assert tb, "Can't get traceback, this shouldn't happen" caller_frame = tb.tb_frame while caller_frame.f_locals.get("self", None) is self: caller_frame = caller_frame.f_back return caller_frame.f_locals def __enter__(self) -> "DeviceCollector": # Stash the names that were defined before we were called self._names_on_enter = set(self._caller_locals()) return self async def __aenter__(self) -> "DeviceCollector": return self.__enter__() async def _on_exit(self) -> None: # Name and kick off connect for devices connect_coroutines: Dict[str, Coroutine] = {} for name, obj in self._objects_on_exit.items(): if name not in self._names_on_enter and isinstance(obj, Device): if self._set_name and not obj.name: obj.set_name(name) if self._connect: connect_coroutines[name] = obj.connect( self._sim, timeout=self._timeout ) # Connect to all the devices if connect_coroutines: await wait_for_connection(**connect_coroutines) async def __aexit__(self, type, value, traceback): self._objects_on_exit = self._caller_locals() await self._on_exit() def __exit__(self, type_, value, traceback): self._objects_on_exit = self._caller_locals() try: fut = call_in_bluesky_event_loop(self._on_exit()) except RuntimeError: raise NotConnected( "Could not connect devices. Is the bluesky event loop running? See " "https://blueskyproject.io/ophyd-async/main/" "user/explanations/event-loop-choice.html for more info." ) return fut