from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Generic, TypeVar
from tango import DevFailed as TangoDevFailed
from tango import DeviceProxy
from tango.asyncio import DeviceProxy as AsyncDeviceProxy
from ophyd_async.core import (
Command,
Device,
DeviceConnector,
DeviceFiller,
DeviceVector,
LazyMock,
NotConnectedError,
Signal,
)
from ._signal import TangoSignalBackend, infer_python_type, infer_signal_type
from ._tango_command_backend import TangoCommandBackend
from ._utils import get_full_attr_trl, sig_from_types
T = TypeVar("T")
[docs]
class TangoDevice(Device):
"""General class for TangoDevices.
Extends Device to provide attributes for Tango devices.
:param trl: Tango resource locator, typically of the device server.
An asynchronous DeviceProxy object will be created using the
trl and awaited when the device is connected.
:param connector: A pre-built `DeviceConnector`, used when this device is
created as a declarative sub-device of another `TangoDevice` (the parent
connector builds it). Mutually exclusive with `trl`: pass one or the
other, not both. When given, `support_events`/`auto_fill_signals` are
ignored (they are configured on the supplied connector instead).
"""
_trl: str = ""
def __init__(
self,
trl: str = "",
support_events: bool = False,
name: str = "",
auto_fill_signals: bool = True,
connector: DeviceConnector | None = None,
) -> None:
if connector is not None and trl:
raise ValueError("TangoDevice takes either `trl` or `connector`, not both")
self._trl = trl
if connector is None:
connector = TangoDeviceConnector(
trl=trl,
support_events=support_events,
auto_fill_signals=auto_fill_signals,
)
super().__init__(name=name, connector=connector)
[docs]
def get_trl(self) -> str:
return self._trl
[docs]
def get_proxy(self) -> DeviceProxy | None:
connector = self._connector
if isinstance(connector, TangoDeviceConnector):
return connector.proxy
return None
[docs]
@dataclass
class TangoPolling(Generic[T]):
ophyd_polling_period: float = 0.1
abs_change: T | None = None
rel_change: T | None = None
def fill_backend_with_polling(
support_events: bool, backend: TangoSignalBackend, annotations: list[Any]
):
unhandled = []
while annotations:
annotation = annotations.pop(0)
backend.allow_events(support_events)
if isinstance(annotation, TangoPolling):
backend.set_polling(
not support_events,
annotation.ophyd_polling_period,
annotation.abs_change,
annotation.rel_change,
)
else:
unhandled.append(annotation)
annotations.extend(unhandled)
# These leftover annotations will now be handled by the iterator
[docs]
class TangoDeviceConnector(DeviceConnector):
def __init__(
self,
trl: str | None,
support_events: bool,
auto_fill_signals: bool = True,
) -> None:
self.trl = trl
self.proxy: DeviceProxy | None = None
self._support_events = support_events
self._auto_fill_signals = auto_fill_signals
[docs]
def set_trl(self, trl: str):
self.trl = trl
[docs]
def create_children_from_annotations(self, device: Device):
if not hasattr(self, "filler"):
self.filler = DeviceFiller(
device=device,
signal_backend_factory=TangoSignalBackend,
device_connector_factory=lambda: TangoDeviceConnector(
None, self._support_events
),
command_backend_factory=TangoCommandBackend,
)
list(self.filler.create_devices_from_annotations(filled=False))
for backend, annotations in self.filler.create_signals_from_annotations(
filled=False
):
fill_backend_with_polling(self._support_events, backend, annotations)
list(self.filler.create_commands_from_annotations(filled=False))
self.filler.check_created()
[docs]
async def connect_mock(self, device: Device, mock: LazyMock):
if isinstance(device, DeviceVector):
# Make 2 entries for this DeviceVector
self.filler.create_device_dict_entries_to_mock([1, 2])
# Set the name of the device to name all children
device.set_name(device.name)
return await super().connect_mock(device, mock)
[docs]
async def connect_real(self, device: Device, timeout: float, force_reconnect: bool):
if not self.trl:
raise RuntimeError(f"Could not created Device Proxy for TRL {self.trl}")
try:
proxy = await AsyncDeviceProxy(self.trl) # type: ignore
except TangoDevFailed as ex:
raise NotConnectedError(f"tango://{self.trl}") from ex
self.proxy = proxy
children = sorted(
set().union(proxy.get_attribute_list()).union(proxy.get_command_list())
) # type: ignore
children = [
child for child in children if child not in self.filler.ignored_signals
]
not_filled = {unfilled for unfilled, _ in device.children()}
# If auto_fill_signals is True, fill all children inferred from the device
# else fill only the children that are annotated
for name in children:
if self._auto_fill_signals or name in not_filled:
# TODO: strip attribute name
full_trl = get_full_attr_trl(self.trl, name)
signal_type = await infer_signal_type(full_trl, proxy)
if signal_type is not None and issubclass(signal_type, Signal):
backend = self.filler.fill_child_signal(name, signal_type)
elif signal_type is not None and issubclass(signal_type, Command):
backend = self.filler.fill_child_command(name, signal_type)
else:
raise TypeError(
f"Cannot infer type for {full_trl} (type {signal_type})"
)
# don't overload datatype if provided by annotation
if isinstance(backend, TangoCommandBackend):
if backend.signature is None:
in_type, out_type = await infer_python_type(full_trl, proxy)
sig = sig_from_types(in_type, out_type)
# # Pyright still thinks backend could be a
# # TangoSignalBackend somehow
backend.signature = sig # type: ignore
elif isinstance(backend, TangoSignalBackend):
if backend.datatype is None:
_, out_type = await infer_python_type(full_trl, proxy)
backend.datatype = out_type
else:
raise TypeError(f"Unknown backend type {backend}")
backend.set_trl(full_trl)
# Check that all the requested children have been filled
self.filler.check_filled(f"{self.trl}: {children}")
# Set the name of the device to name all children
device.set_name(device.name)
return await super().connect_real(device, timeout, force_reconnect)
