ophyd-async 0.3a1__py3-none-any.whl → 0.3a3__py3-none-any.whl

ophyd_async/epics/pvi/pvi.py

@@ -0,0 +1,318 @@
+import re
+from dataclasses import dataclass
+from inspect import isclass
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    FrozenSet,
+    Literal,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+from ophyd_async.core import Device, DeviceVector, SimSignalBackend
+from ophyd_async.core.signal import Signal
+from ophyd_async.core.utils import DEFAULT_TIMEOUT
+from ophyd_async.epics._backend._p4p import PvaSignalBackend
+from ophyd_async.epics.signal.signal import (
+    epics_signal_r,
+    epics_signal_rw,
+    epics_signal_w,
+    epics_signal_x,
+)
+
+T = TypeVar("T")
+Access = FrozenSet[
+    Union[Literal["r"], Literal["w"], Literal["rw"], Literal["x"], Literal["d"]]
+]
+
+
+def _strip_number_from_string(string: str) -> Tuple[str, Optional[int]]:
+    match = re.match(r"(.*?)(\d*)$", string)
+    assert match
+
+    name = match.group(1)
+    number = match.group(2) or None
+    if number:
+        number = int(number)
+    return name, number
+
+
+def _split_subscript(tp: T) -> Union[Tuple[Any, Tuple[Any]], Tuple[T, None]]:
+    """Split a subscripted type into the its origin and args.
+
+    If `tp` is not a subscripted type, then just return the type and None as args.
+
+    """
+    if get_origin(tp) is not None:
+        return get_origin(tp), get_args(tp)
+
+    return tp, None
+
+
+def _strip_union(field: Union[Union[T], T]) -> Tuple[T, bool]:
+    if get_origin(field) is Union:
+        args = get_args(field)
+        is_optional = type(None) in args
+        for arg in args:
+            if arg is not type(None):
+                return arg, is_optional
+    return field, False
+
+
+def _strip_device_vector(field: Union[Type[Device]]) -> Tuple[bool, Type[Device]]:
+    if get_origin(field) is DeviceVector:
+        return True, get_args(field)[0]
+    return False, field
+
+
+@dataclass
+class PVIEntry:
+    """
+    A dataclass to represent a single entry in the PVI table.
+    This could either be a signal or a sub-table.
+    """
+
+    sub_entries: Dict[str, Union[Dict[int, "PVIEntry"], "PVIEntry"]]
+    pvi_pv: Optional[str] = None
+    device: Optional[Device] = None
+    common_device_type: Optional[Type[Device]] = None
+
+
+def _verify_common_blocks(entry: PVIEntry, common_device: Type[Device]):
+    if not entry.sub_entries:
+        return
+    common_sub_devices = get_type_hints(common_device)
+    for sub_name, sub_device in common_sub_devices.items():
+        if sub_name in ("_name", "parent"):
+            continue
+        assert entry.sub_entries
+        device_t, is_optional = _strip_union(sub_device)
+        if sub_name not in entry.sub_entries and not is_optional:
+            raise RuntimeError(
+                f"sub device `{sub_name}:{type(sub_device)}` " "was not provided by pvi"
+            )
+        if isinstance(entry.sub_entries[sub_name], dict):
+            for sub_sub_entry in entry.sub_entries[sub_name].values():  # type: ignore
+                _verify_common_blocks(sub_sub_entry, sub_device)  # type: ignore
+        else:
+            _verify_common_blocks(
+                entry.sub_entries[sub_name],
+                sub_device,  # type: ignore
+            )
+
+
+_pvi_mapping: Dict[FrozenSet[str], Callable[..., Signal]] = {
+    frozenset({"r", "w"}): lambda dtype, read_pv, write_pv: epics_signal_rw(
+        dtype, "pva://" + read_pv, "pva://" + write_pv
+    ),
+    frozenset({"rw"}): lambda dtype, read_write_pv: epics_signal_rw(
+        dtype, "pva://" + read_write_pv, write_pv="pva://" + read_write_pv
+    ),
+    frozenset({"r"}): lambda dtype, read_pv: epics_signal_r(dtype, "pva://" + read_pv),
+    frozenset({"w"}): lambda dtype, write_pv: epics_signal_w(
+        dtype, "pva://" + write_pv
+    ),
+    frozenset({"x"}): lambda _, write_pv: epics_signal_x("pva://" + write_pv),
+}
+
+
+def _parse_type(
+    is_pvi_table: bool,
+    number_suffix: Optional[int],
+    common_device_type: Optional[Type[Device]],
+):
+    if common_device_type:
+        # pre-defined type
+        device_cls, _ = _strip_union(common_device_type)
+        is_device_vector, device_cls = _strip_device_vector(device_cls)
+        device_cls, device_args = _split_subscript(device_cls)
+        assert issubclass(device_cls, Device)
+
+        is_signal = issubclass(device_cls, Signal)
+        signal_dtype = device_args[0] if device_args is not None else None
+
+    elif is_pvi_table:
+        # is a block, we can make it a DeviceVector if it ends in a number
+        is_device_vector = number_suffix is not None
+        is_signal = False
+        signal_dtype = None
+        device_cls = Device
+    else:
+        # is a signal, signals aren't stored in DeviceVectors unless
+        # they're defined as such in the common_device_type
+        is_device_vector = False
+        is_signal = True
+        signal_dtype = None
+        device_cls = Signal
+
+    return is_device_vector, is_signal, signal_dtype, device_cls
+
+
+def _sim_common_blocks(device: Device, stripped_type: Optional[Type] = None):
+    device_t = stripped_type or type(device)
+    sub_devices = (
+        (field, field_type)
+        for field, field_type in get_type_hints(device_t).items()
+        if field not in ("_name", "parent")
+    )
+
+    for device_name, device_cls in sub_devices:
+        device_cls, _ = _strip_union(device_cls)
+        is_device_vector, device_cls = _strip_device_vector(device_cls)
+        device_cls, device_args = _split_subscript(device_cls)
+        assert issubclass(device_cls, Device)
+
+        is_signal = issubclass(device_cls, Signal)
+        signal_dtype = device_args[0] if device_args is not None else None
+
+        if is_device_vector:
+            if is_signal:
+                sub_device_1 = device_cls(SimSignalBackend(signal_dtype))
+                sub_device_2 = device_cls(SimSignalBackend(signal_dtype))
+                sub_device = DeviceVector({1: sub_device_1, 2: sub_device_2})
+            else:
+                sub_device = DeviceVector({1: device_cls(), 2: device_cls()})
+
+                for sub_device_in_vector in sub_device.values():
+                    _sim_common_blocks(sub_device_in_vector, stripped_type=device_cls)
+
+            for value in sub_device.values():
+                value.parent = sub_device
+        else:
+            if is_signal:
+                sub_device = device_cls(SimSignalBackend(signal_dtype))
+            else:
+                sub_device = getattr(device, device_name, device_cls())
+                _sim_common_blocks(sub_device, stripped_type=device_cls)
+
+        setattr(device, device_name, sub_device)
+        sub_device.parent = device
+
+
+async def _get_pvi_entries(entry: PVIEntry, timeout=DEFAULT_TIMEOUT):
+    if not entry.pvi_pv or not entry.pvi_pv.endswith(":PVI"):
+        raise RuntimeError("Top level entry must be a pvi table")
+
+    pvi_table_signal_backend: PvaSignalBackend = PvaSignalBackend(
+        None, entry.pvi_pv, entry.pvi_pv
+    )
+    await pvi_table_signal_backend.connect(
+        timeout=timeout
+    )  # create table signal backend
+
+    pva_table = (await pvi_table_signal_backend.get_value())["pvi"]
+    common_device_type_hints = (
+        get_type_hints(entry.common_device_type) if entry.common_device_type else {}
+    )
+
+    for sub_name, pva_entries in pva_table.items():
+        pvs = list(pva_entries.values())
+        is_pvi_table = len(pvs) == 1 and pvs[0].endswith(":PVI")
+        sub_name_split, sub_number_split = _strip_number_from_string(sub_name)
+        is_device_vector, is_signal, signal_dtype, device_type = _parse_type(
+            is_pvi_table,
+            sub_number_split,
+            common_device_type_hints.get(sub_name_split),
+        )
+        if is_signal:
+            device = _pvi_mapping[frozenset(pva_entries.keys())](signal_dtype, *pvs)
+        else:
+            device = getattr(entry.device, sub_name, device_type())
+
+        sub_entry = PVIEntry(
+            device=device, common_device_type=device_type, sub_entries={}
+        )
+
+        if is_device_vector:
+            # If device vector then we store sub_name -> {sub_number -> sub_entry}
+            # and aggregate into `DeviceVector` in `_set_device_attributes`
+            sub_number_split = 1 if sub_number_split is None else sub_number_split
+            if sub_name_split not in entry.sub_entries:
+                entry.sub_entries[sub_name_split] = {}
+            entry.sub_entries[sub_name_split][sub_number_split] = sub_entry  # type: ignore
+        else:
+            entry.sub_entries[sub_name] = sub_entry
+
+        if is_pvi_table:
+            sub_entry.pvi_pv = pvs[0]
+            await _get_pvi_entries(sub_entry)
+
+    if entry.common_device_type:
+        _verify_common_blocks(entry, entry.common_device_type)
+
+
+def _set_device_attributes(entry: PVIEntry):
+    for sub_name, sub_entry in entry.sub_entries.items():
+        if isinstance(sub_entry, dict):
+            sub_device = DeviceVector()  # type: ignore
+            for key, device_vector_sub_entry in sub_entry.items():
+                sub_device[key] = device_vector_sub_entry.device
+                if device_vector_sub_entry.pvi_pv:
+                    _set_device_attributes(device_vector_sub_entry)
+                # Set the device vector entry to have the device vector as a parent
+                device_vector_sub_entry.device.parent = sub_device  # type: ignore
+        else:
+            sub_device = sub_entry.device  # type: ignore
+            if sub_entry.pvi_pv:
+                _set_device_attributes(sub_entry)
+
+        sub_device.parent = entry.device
+        setattr(entry.device, sub_name, sub_device)
+
+
+async def fill_pvi_entries(
+    device: Device, root_pv: str, timeout=DEFAULT_TIMEOUT, sim=False
+):
+    """
+    Fills a ``device`` with signals from a the ``root_pvi:PVI`` table.
+
+    If the device names match with parent devices of ``device`` then types are used.
+    """
+    if sim:
+        # set up sim signals for the common annotations
+        _sim_common_blocks(device)
+    else:
+        # check the pvi table for devices and fill the device with them
+        root_entry = PVIEntry(
+            pvi_pv=root_pv,
+            device=device,
+            common_device_type=type(device),
+            sub_entries={},
+        )
+        await _get_pvi_entries(root_entry, timeout=timeout)
+        _set_device_attributes(root_entry)
+
+    # We call set name now the parent field has been set in all of the
+    # introspect-initialized devices. This will recursively set the names.
+    device.set_name(device.name)
+
+
+def create_children_from_annotations(
+    device: Device, included_optional_fields: Tuple[str, ...] = ()
+):
+    """For intializing blocks at __init__ of ``device``."""
+    for name, device_type in get_type_hints(type(device)).items():
+        if name in ("_name", "parent"):
+            continue
+        device_type, is_optional = _strip_union(device_type)
+        if is_optional and name not in included_optional_fields:
+            continue
+        is_device_vector, device_type = _strip_device_vector(device_type)
+        if (
+            is_device_vector
+            or ((origin := get_origin(device_type)) and issubclass(origin, Signal))
+            or (isclass(device_type) and issubclass(device_type, Signal))
+        ):
+            continue
+
+        sub_device = device_type()
+        setattr(device, name, sub_device)
+        create_children_from_annotations(sub_device)

ophyd_async/epics/signal/signal.py

@@ -41,7 +41,7 @@ def _make_backend(
 
 
 def epics_signal_rw(
-    datatype: Type[T], read_pv: str, write_pv: Optional[str] = None
+    datatype: Type[T], read_pv: str, write_pv: Optional[str] = None, name: str = ""
 ) -> SignalRW[T]:
     """Create a `SignalRW` backed by 1 or 2 EPICS PVs
 
@@ -55,24 +55,41 @@ def epics_signal_rw(
         If given, use this PV to write to, otherwise use read_pv
     """
     backend = _make_backend(datatype, read_pv, write_pv or read_pv)
-    return SignalRW(backend)
+    return SignalRW(backend, name=name)
 
 
-def epics_signal_r(datatype: Type[T], read_pv: str) -> SignalR[T]:
-    """Create a `SignalR` backed by 1 EPICS PV
+def epics_signal_rw_rbv(
+    datatype: Type[T], write_pv: str, read_suffix: str = "_RBV", name: str = ""
+) -> SignalRW[T]:
+    """Create a `SignalRW` backed by 1 or 2 EPICS PVs, with a suffix on the readback pv
 
     Parameters
     ----------
     datatype:
         Check that the PV is of this type
+    write_pv:
+        The PV to write to
+    read_suffix:
+        Append this suffix to the write pv to create the readback pv
+    """
+    return epics_signal_rw(datatype, f"{write_pv}{read_suffix}", write_pv, name)
+
+
+def epics_signal_r(datatype: Type[T], read_pv: str, name: str = "") -> SignalR[T]:
+    """Create a `SignalR` backed by 1 EPICS PV
+
+    Parameters
+    ----------
+    datatype
+        Check that the PV is of this type
     read_pv:
         The PV to read and monitor
     """
     backend = _make_backend(datatype, read_pv, read_pv)
-    return SignalR(backend)
+    return SignalR(backend, name=name)
 
 
-def epics_signal_w(datatype: Type[T], write_pv: str) -> SignalW[T]:
+def epics_signal_w(datatype: Type[T], write_pv: str, name: str = "") -> SignalW[T]:
     """Create a `SignalW` backed by 1 EPICS PVs
 
     Parameters
@@ -83,10 +100,10 @@ def epics_signal_w(datatype: Type[T], write_pv: str) -> SignalW[T]:
         The PV to write to
     """
     backend = _make_backend(datatype, write_pv, write_pv)
-    return SignalW(backend)
+    return SignalW(backend, name=name)
 
 
-def epics_signal_x(write_pv: str) -> SignalX:
+def epics_signal_x(write_pv: str, name: str = "") -> SignalX:
     """Create a `SignalX` backed by 1 EPICS PVs
 
     Parameters
@@ -95,4 +112,4 @@ def epics_signal_x(write_pv: str) -> SignalX:
         The PV to write its initial value to on trigger
     """
     backend: SignalBackend = _make_backend(None, write_pv, write_pv)
-    return SignalX(backend)
+    return SignalX(backend, name=name)

ophyd_async/log.py

@@ -0,0 +1,130 @@
+import logging
+import sys
+
+import colorlog
+
+__all__ = (
+    "config_ophyd_async_logging",
+    "logger",
+    "set_handler",
+)
+
+DEFAULT_FORMAT = (
+    "%(log_color)s[%(levelname)1.1s %(asctime)s.%(msecs)03d "
+    "%(module)s:%(lineno)d] %(message)s"
+)
+
+DEFAULT_DATE_FORMAT = "%y%m%d %H:%M:%S"
+
+DEFAULT_LOG_COLORS = {
+    "DEBUG": "cyan",
+    "INFO": "green",
+    "WARNING": "yellow",
+    "ERROR": "red",
+    "CRITICAL": "red,bg_white",
+}
+
+
+class ColoredFormatterWithDeviceName(colorlog.ColoredFormatter):
+    def format(self, record):
+        message = super().format(record)
+        if hasattr(record, "ophyd_async_device_name"):
+            message = f"[{record.ophyd_async_device_name}]{message}"
+        return message
+
+
+def _validate_level(level) -> int:
+    """
+    Return an int for level comparison
+    """
+    if isinstance(level, int):
+        levelno = level
+    elif isinstance(level, str):
+        levelno = logging.getLevelName(level)
+
+    if isinstance(levelno, int):
+        return levelno
+    else:
+        raise ValueError(
+            "Your level is illegal, please use "
+            "'CRITICAL', 'FATAL', 'ERROR', 'WARNING', 'INFO', or 'DEBUG'."
+        )
+
+
+logger = logging.getLogger("ophyd_async")
+
+current_handler = None  # overwritten below
+
+
+def config_ophyd_async_logging(
+    file=sys.stdout,
+    fmt=DEFAULT_FORMAT,
+    datefmt=DEFAULT_DATE_FORMAT,
+    color=True,
+    level="WARNING",
+):
+    """
+    Set a new handler on the ``logging.getLogger('ophyd_async')`` logger.
+    If this is called more than once, the handler from the previous invocation
+    is removed (if still present) and replaced.
+
+    Parameters
+    ----------
+    file : object with ``write`` method or filename string
+        Default is ``sys.stdout``.
+    fmt : Overall logging format
+    datefmt : string
+        Date format. Default is ``'%H:%M:%S'``.
+    color : boolean
+        Use ANSI color codes. True by default.
+    level : str or int
+        Python logging level, given as string or corresponding integer.
+        Default is 'WARNING'.
+    Returns
+    -------
+    handler : logging.Handler
+        The handler, which has already been added to the 'ophyd_async' logger.
+    Examples
+    --------
+    Log to a file.
+        config_ophyd_async_logging(file='/tmp/what_is_happening.txt')
+    Include the date along with the time. (The log messages will always include
+    microseconds, which are configured separately, not as part of 'datefmt'.)
+        config_ophyd_async_logging(datefmt="%Y-%m-%d %H:%M:%S")
+    Turn off ANSI color codes.
+        config_ophyd_async_logging(color=False)
+    Increase verbosity: show level DEBUG or higher.
+        config_ophyd_async_logging(level='DEBUG')
+    """
+    global current_handler
+
+    if isinstance(file, str):
+        handler = logging.FileHandler(file)
+        formatter = ColoredFormatterWithDeviceName(
+            fmt=fmt, datefmt=datefmt, no_color=True
+        )
+    else:
+        handler = colorlog.StreamHandler(file)
+        formatter = ColoredFormatterWithDeviceName(
+            fmt=fmt, datefmt=datefmt, log_colors=DEFAULT_LOG_COLORS, no_color=color
+        )
+
+    levelno = _validate_level(level)
+    handler.setFormatter(formatter)
+    handler.setLevel(levelno)
+
+    if current_handler in logger.handlers:
+        logger.removeHandler(current_handler)
+    logger.addHandler(handler)
+
+    current_handler = handler
+
+    if logger.getEffectiveLevel() > levelno:
+        logger.setLevel(levelno)
+    try:
+        return handler
+    finally:
+        handler.close()
+
+
+set_handler = config_ophyd_async_logging  # for back-compat

ophyd_async/panda/__init__.py

@@ -1,19 +1,27 @@
-from .panda import PandA, PcapBlock, PulseBlock, PVIEntry, SeqBlock, SeqTable
-from .panda_controller import PandaPcapController
-from .table import (
+from ._common_blocks import (
+    CommonPandaBlocks,
+    DataBlock,
+    PcapBlock,
+    PulseBlock,
+    SeqBlock,
+    TimeUnits,
+)
+from ._hdf_panda import HDFPanda
+from ._panda_controller import PandaPcapController
+from ._table import (
     SeqTable,
     SeqTableRow,
     SeqTrigger,
     seq_table_from_arrays,
     seq_table_from_rows,
 )
-from .utils import phase_sorter
+from ._utils import phase_sorter
 
 __all__ = [
-    "PandA",
+    "CommonPandaBlocks",
+    "HDFPanda",
     "PcapBlock",
     "PulseBlock",
-    "PVIEntry",
     "seq_table_from_arrays",
     "seq_table_from_rows",
     "SeqBlock",
@@ -22,4 +30,7 @@ __all__ = [
     "SeqTrigger",
     "phase_sorter",
     "PandaPcapController",
+    "TimeUnits",
+    "DataBlock",
+    "CommonPandABlocks",
 ]

ophyd_async/panda/_common_blocks.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from enum import Enum
+
+from ophyd_async.core import Device, DeviceVector, SignalR, SignalRW
+from ophyd_async.panda._table import SeqTable
+
+
+class DataBlock(Device):
+    # In future we may decide to make hdf_* optional
+    hdf_directory: SignalRW[str]
+    hdf_file_name: SignalRW[str]
+    num_capture: SignalRW[int]
+    num_captured: SignalR[int]
+    capture: SignalRW[bool]
+    flush_period: SignalRW[float]
+
+
+class PulseBlock(Device):
+    delay: SignalRW[float]
+    width: SignalRW[float]
+
+
+class TimeUnits(str, Enum):
+    min = "min"
+    s = "s"
+    ms = "ms"
+    us = "us"
+
+
+class SeqBlock(Device):
+    table: SignalRW[SeqTable]
+    active: SignalRW[bool]
+    repeats: SignalRW[int]
+    prescale: SignalRW[float]
+    prescale_units: SignalRW[TimeUnits]
+    enable: SignalRW[str]
+
+
+class PcapBlock(Device):
+    active: SignalR[bool]
+    arm: SignalRW[bool]
+
+
+class CommonPandaBlocks(Device):
+    pulse: DeviceVector[PulseBlock]
+    seq: DeviceVector[SeqBlock]
+    pcap: PcapBlock
+    data: DataBlock

ophyd_async/panda/_hdf_panda.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from typing import Sequence
+
+from ophyd_async.core import (
+    DEFAULT_TIMEOUT,
+    DirectoryProvider,
+    SignalR,
+    StandardDetector,
+)
+from ophyd_async.epics.pvi import create_children_from_annotations, fill_pvi_entries
+
+from ._common_blocks import CommonPandaBlocks
+from ._panda_controller import PandaPcapController
+from .writers._hdf_writer import PandaHDFWriter
+
+
+class HDFPanda(CommonPandaBlocks, StandardDetector):
+    def __init__(
+        self,
+        prefix: str,
+        directory_provider: DirectoryProvider,
+        config_sigs: Sequence[SignalR] = (),
+        name: str = "",
+    ):
+        self._prefix = prefix
+
+        create_children_from_annotations(self)
+        controller = PandaPcapController(pcap=self.pcap)
+        writer = PandaHDFWriter(
+            prefix=prefix,
+            directory_provider=directory_provider,
+            name_provider=lambda: name,
+            panda_device=self,
+        )
+        super().__init__(
+            controller=controller,
+            writer=writer,
+            config_sigs=config_sigs,
+            name=name,
+            writer_timeout=DEFAULT_TIMEOUT,
+        )
+
+    async def connect(
+        self, sim: bool = False, timeout: float = DEFAULT_TIMEOUT
+    ) -> None:
+        await fill_pvi_entries(self, self._prefix + "PVI", timeout=timeout, sim=sim)
+        await super().connect(sim=sim, timeout=timeout)

ophyd_async/panda/{panda_controller.py → _panda_controller.py}

@@ -7,15 +7,11 @@ from ophyd_async.core import (
     DetectorTrigger,
     wait_for_value,
 )
-
-from .panda import PcapBlock
+from ophyd_async.panda import PcapBlock
 
 
 class PandaPcapController(DetectorControl):
-    def __init__(
-        self,
-        pcap: PcapBlock,
-    ) -> None:
+    def __init__(self, pcap: PcapBlock) -> None:
         self.pcap = pcap
 
     def get_deadtime(self, exposure: float) -> float:
@@ -35,7 +31,7 @@ class PandaPcapController(DetectorControl):
         await wait_for_value(self.pcap.active, True, timeout=1)
         return AsyncStatus(wait_for_value(self.pcap.active, False, timeout=None))
 
-    async def disarm(self):
+    async def disarm(self) -> AsyncStatus:
         await asyncio.gather(self.pcap.arm.set(False))
         await wait_for_value(self.pcap.active, False, timeout=1)
         return AsyncStatus(wait_for_value(self.pcap.active, False, timeout=None))