Fix DeviceVector introspection

src/ophyd_async/epics/core/__init__.py

@@ -1,6 +1,6 @@
 from ._epics_connector import EpicsDeviceConnector, PvSuffix
 from ._epics_device import EpicsDevice
-from ._pvi_connector import PviDeviceConnector
+from ._pvi_connector import PviDeviceConnector, PviTree, SignalDetails
 from ._signal import (
     CaSignalBackend,
     PvaSignalBackend,
@@ -14,6 +14,8 @@
 
 __all__ = [
     "PviDeviceConnector",
+    "PviTree",
+    "SignalDetails",
     "EpicsDeviceConnector",
     "PvSuffix",
     "EpicsDevice",

src/ophyd_async/epics/core/_pvi_connector.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
-from typing import Literal, cast
+import re
+
+from pydantic import Field
 
 from ophyd_async.core import (
+    ConfinedModel,
     Device,
     DeviceConnector,
     DeviceFiller,
@@ -12,56 +15,12 @@
     SignalRW,
     SignalW,
     SignalX,
+    gather_dict,
 )
 
 from ._epics_connector import fill_backend_with_prefix
 from ._signal import PvaSignalBackend, pvget_with_timeout
 
-Entry = dict[str, str]
-
-OldPVIVector = list[Entry | None]
-# The older PVI structure has vectors of the form
-# structure[] ttlout
-#     (none)
-#     structure
-#         string d PANDABLOCKS_IOC:TTLOUT1:PVI
-#     structure
-#         string d PANDABLOCKS_IOC:TTLOUT2:PVI
-#     structure
-#         string d PANDABLOCKS_IOC:TTLOUT3:PVI
-
-
-FastCSPVIVector = dict[Literal["d"], Entry]
-# The newer pva FastCS PVI structure has vectors of the form
-# structure ttlout
-#     structure d
-#         string v1 FASTCS_PANDA:Ttlout1:PVI
-#         string v2 FASTCS_PANDA:Ttlout2:PVI
-#         string v3 FASTCS_PANDA:Ttlout3:PVI
-#         string v4 FASTCS_PANDA:Ttlout4:PVI
-
-
-def _get_signal_details(entry: Entry) -> tuple[type[Signal], str, str]:
-    match entry:
-        case {"r": read_pv, "w": write_pv}:
-            return SignalRW, read_pv, write_pv
-        case {"r": read_pv}:
-            return SignalR, read_pv, read_pv
-        case {"w": write_pv}:
-            return SignalW, write_pv, write_pv
-        case {"rw": read_write_pv}:
-            return SignalRW, read_write_pv, read_write_pv
-        case {"x": execute_pv}:
-            return SignalX, execute_pv, execute_pv
-        case _:
-            raise TypeError(f"Can't process entry {entry}")
-
-
-def _is_device_vector_entry(entry: Entry | OldPVIVector | FastCSPVIVector) -> bool:
-    return isinstance(entry, list) or (
-        entry.keys() == {"d"} and isinstance(entry["d"], dict)
-    )
-
 
 class PviDeviceConnector(DeviceConnector):
     """Connect to PVI structure served over PVA.
@@ -78,6 +37,7 @@
     """
 
     mock_device_vector_len: int = 2
+    pvi_tree: PviTree | None = None
 
     def __init__(self, prefix: str = "", error_hint: str = "") -> None:
         # TODO: what happens if we get a leading "pva://" here?
@@ -101,52 +61,215 @@
                 fill_backend_with_prefix(self.prefix, backend, annotations)
             self.filler.check_created()
 
-    def _fill_child(self, name: str, entry: Entry, vector_index: int | None = None):
-        if set(entry) == {"d"}:
-            connector = self.filler.fill_child_device(name, vector_index=vector_index)
-            connector.pvi_pv = entry["d"]
-        else:
-            signal_type, read_pv, write_pv = _get_signal_details(entry)
-            backend = self.filler.fill_child_signal(name, signal_type, vector_index)
-            backend.read_pv = read_pv
-            backend.write_pv = write_pv
-
     async def connect_mock(self, device: Device, mock: LazyMock):
         self.filler.create_device_vector_entries_to_mock(self.mock_device_vector_len)
         # Set the name of the device to name all children
         device.set_name(device.name)
         return await super().connect_mock(device, mock)
 
-    def _fill_vector_child(self, name: str, entry: OldPVIVector | FastCSPVIVector):
-        if isinstance(entry, list):
-            for i, e in enumerate(entry):
-                if e:
-                    self._fill_child(name, e, i)
-        else:
-            for i_string, e in entry["d"].items():
-                self._fill_child(name, {"d": e}, int(i_string.lstrip("v")))
-
     async def connect_real(
         self, device: Device, timeout: float, force_reconnect: bool
     ) -> None:
-        pvi_structure = await pvget_with_timeout(self.pvi_pv, timeout)
-
-        entries: dict[str, Entry | OldPVIVector | FastCSPVIVector] = pvi_structure[
-            "value"
-        ].todict()
-        # Fill based on what PVI gives us
-        for name, entry in entries.items():
-            if _is_device_vector_entry(entry):
-                self._fill_vector_child(
-                    name, cast(OldPVIVector | FastCSPVIVector, entry)
-                )
+        if not self.pvi_tree:
+            # Top-level device, so discover PVI tree
+            self.pvi_tree = await PviTree.build_device_tree(
+                name=device.name, pvi_pv=self.pvi_pv, timeout=timeout
+            )
+        # Fill all signals
+        for signal_name, signal_details in self.pvi_tree.signals.items():
+            backend = self.filler.fill_child_signal(
+                signal_name, signal_details.signal_type, None
+            )
+            backend.read_pv = signal_details.read_pv
+            backend.write_pv = signal_details.write_pv
+        # Fill all sub devices
+        for device_name, device_sub_tree in self.pvi_tree.sub_devices.items():
+            if device_sub_tree.vector_children:
+                # This is a DeviceVector
+                for (
+                    vector_index,
+                    vector_child,
+                ) in device_sub_tree.vector_children.items():
+                    connector = self.filler.fill_child_device(
+                        device_name, vector_index=vector_index
+                    )
+                    connector.pvi_tree = vector_child
+                    connector.pvi_pv = vector_child.pvi_pv
             else:
-                # This is a child
-                self._fill_child(name, cast(Entry, entry))
+                # This is a Device
+                connector = self.filler.fill_child_device(device_name)
+                connector.pvi_tree = device_sub_tree
+                connector.pvi_pv = device_sub_tree.pvi_pv
 
         # Check that all the requested children have been filled
         suffix = f"\n{self.error_hint}" if self.error_hint else ""
-        self.filler.check_filled(f"{self.pvi_pv}: {entries}{suffix}")
+        self.filler.check_filled(f"{self.pvi_pv}: {self.pvi_tree}{suffix}")
         # Set the name of the device to name all children
         device.set_name(device.name)
         return await super().connect_real(device, timeout, force_reconnect)
+
+
+class SignalDetails(ConfinedModel):
+    """Representation of a Signal to be constructed."""
+
+    signal_type: type[Signal]
+    read_pv: str
+    write_pv: str
+
+    @classmethod
+    def from_entry(cls, entry: dict[str, str]) -> SignalDetails:
+        match entry:
+            case {"r": read_pv, "w": write_pv}:
+                return cls(signal_type=SignalRW, read_pv=read_pv, write_pv=write_pv)
+
+            case {"rw": pv}:
+                return cls(signal_type=SignalRW, read_pv=pv, write_pv=pv)
+
+            case {"r": read_pv}:
+                return cls(signal_type=SignalR, read_pv=read_pv, write_pv=read_pv)
+
+            case {"w": write_pv}:
+                return cls(signal_type=SignalW, read_pv=write_pv, write_pv=write_pv)
+
+            case {"x": execute_pv}:
+                return cls(signal_type=SignalX, read_pv=execute_pv, write_pv=execute_pv)
+
+            case _:
+                raise TypeError(f"Can't process entry {entry}")
+
+
+class PviTree(ConfinedModel):
+    """Representation of a PVI structure of devices and signals in a PVI query.
+
+    Example 1: A device with sub-devices and signals
+    --------------------------------------
+    For a PVI structure such as:
+
+    ```json
+    {
+        "bit": {"d": "TEST-PANDA:Bits:PVI"},
+        "calc": {"d": "TEST-PANDA:Calc:PVI"},
+        "a": {"rw": "TEST-PANDA:Bits:A"}
+    }
+    ```
+
+    From "TEST-PANDA:PVI", This would be represented as:
+
+    ```python
+    PviTree(
+        pvi_pv="TEST-PANDA:PVI",
+        signals={
+            "a": SignalDetails(
+                signal_type=SignalRW,
+                read_pv="TEST-PANDA:Bits:A",
+                write_pv="TEST-PANDA:Bits:A")
+        },
+        sub_devices={
+            "bit": PviTree(...),
+            "calc": PviTree(...)
+        },
+        vector_children=[]
+    )
+    ```
+
+    Example 2: A device with vector children
+    -----------------------------------------
+    If an entry like `"calc"` is a **DeviceVector**
+    (e.g., mirroring a fastCS controller vector), the PVI entries will look like this:
+
+    ```json
+    {
+        "__1": {"d": "TEST-PANDA:Calc:2:PVI"},
+        "__2": {"d": "TEST-PANDA:Calc:1:PVI"}
+    }
+    ```
+
+    This would be represented as:
+
+    ```python
+    PviTree(
+        pvi_pv="TEST-PANDA:Calc:PVI",
+        signals={},
+        sub_devices={},
+        vector_children=[
+            PviTree(pvi_pv="TEST-PANDA:Calc:2:PVI", signals={}, ...),
+            PviTree(pvi_pv="TEST-PANDA:Calc:1:PVI", signals={}, ...)
+        ]
+    )
+    ```
+
+    :param pvi_pv:
+        The PVI PV of the device.
+
+    :param signals:
+        A dictionary mapping signal names to `SignalDetails` objects.
+
+    :param sub_devices:
+        A dictionary mapping sub-device names to their corresponding `PviTree` objects.
+
+    :param vector_children:
+        A list of `PviTree` objects representing child devices of a vector device.
+    """
+
+    pvi_pv: str
+    signals: dict[str, SignalDetails] = Field(default={})
+    sub_devices: dict[str, PviTree] = Field(default={})
+    vector_children: dict[int, PviTree] = Field(default={})
+
+    @classmethod
+    async def build_device_tree(cls, name: str, pvi_pv: str, timeout: float) -> PviTree:
+        """Recursively build a PviTree from a top level device.
+
+        Starting from the top-level device, this classmethod performs
+        post-order traversal over the served PVI structure, populating
+        a PviTree from the bottom up.
+
+        :param name: Device name
+        :param pvi_pv: Device PVI PV
+        :param timeout: Timeout on pvget
+        """
+        pvi_structure = await pvget_with_timeout(pvi_pv, timeout)
+
+        # An example entry is: {"d": "Prefix:Device:PVI", "rw": "Prefix:A"}
+        # these entries are stored under the parent PVI structure name
+        # for example, {"device": {"d": "Prefix:Device:PVI", "rw": "Prefix:A"}}
+        entries: dict[str, dict[str, str]] = pvi_structure["value"].todict()
+
+        signal_details = {
+            entry_name: SignalDetails.from_entry(entries.pop(entry_name))
+            for entry_name in entries
+            if set(entries[entry_name]) != {"d"}
+        }
+        sub_trees = await gather_dict(
+            {
+                entry_name: cls.build_device_tree(entry_name, entry["d"], timeout)
+                for entry_name, entry in entries.items()
+            }
+        )
+
+        vector_children: dict[int, PviTree] = {}
+        # Filter vector children out of stand-alone devices
+        for child_name in list(sub_trees):
+            # Check if any sub-devices are named "__#" (e.g., "__1")
+            if m := re.match(r"^__(\d+)$", child_name):
+                sub_tree = sub_trees.pop(child_name)
+                vector_children[int(m.group(1))] = sub_tree
+
+        return PviTree(
+            pvi_pv=pvi_pv,
+            signals=signal_details,
+            sub_devices=sub_trees,
+            vector_children=vector_children,
+        )
+
+    def __str__(self) -> str:
+        """Print a readable top layer of the PviTree."""
+        children = {
+            child_name: tree.pvi_pv
+            for child_name, tree in {**self.sub_devices, **self.vector_children}.items()
+        }
+        signals = {
+            signal_name: (detail.signal_type.__name__, detail.read_pv, detail.write_pv)
+            for signal_name, detail in self.signals.items()
+        }
+        return f"sub_devices={children}\nsignals={signals}"

tests/system_tests/fastcs/panda/test_panda_connect.py

@@ -37,12 +37,14 @@ def __init__(self, uri: str, name: str = ""):
 @pytest.mark.timeout(15.0 if os.name == "nt" else 4.0)
 async def test_panda_with_missing_blocks(panda_pva, panda_t):
     panda = panda_t("PANDAQSRVI:", name="mypanda")
+
     with pytest.raises(
         RuntimeError,
         match=re.escape(
             "mypanda: cannot provision ['pcap'] from PANDAQSRVI:PVI: "
-            "{'pulse': [None, {'d': 'PANDAQSRVI:PULSE1:PVI'}],"
-            " 'seq': [None, {'d': 'PANDAQSRVI:SEQ1:PVI'}]}\nIs it ok?"
+            "sub_devices={'pulse': 'PANDAQSRVI:PULSE:PVI', 'seq': 'PANDAQSRVI:SEQ:PVI'}"
+            "\nsignals={}"
+            "\nIs it ok?"
         ),
     ):
         await panda.connect()

tests/unit_tests/epics/pvi/test_pvi.py

@@ -1,6 +1,5 @@
 from typing import Annotated as A
 from typing import TypeVar
-from unittest.mock import MagicMock
 
 import pytest
 from bluesky.protocols import HasHints, Hints
@@ -16,7 +15,7 @@
     init_devices,
 )
 from ophyd_async.core import StandardReadableFormat as Format
-from ophyd_async.epics.core import PviDeviceConnector
+from ophyd_async.epics.core import PviDeviceConnector, SignalDetails
 
 
 class Block1(Device, HasHints):
@@ -207,14 +206,10 @@ async def test_no_type_annotation_blocks(cls):
 async def test_correctly_setting_signal_type_from_signal_details(
     mock_entry, expected_signal_type
 ):
-    connector = PviDeviceConnector("")
-    connector.filler = MagicMock()
     if not expected_signal_type:
         with pytest.raises(TypeError) as exc:
-            connector._fill_child("signal", mock_entry)
+            SignalDetails.from_entry(mock_entry)
         assert "Can't process entry" in str(exc.value)
     else:
-        connector._fill_child("signal", mock_entry)
-        connector.filler.fill_child_signal.assert_called_once_with(
-            "signal", expected_signal_type, None
-        )
+        details = SignalDetails.from_entry(mock_entry)
+        assert details.signal_type == expected_signal_type