Implement SubControllerVector (#192)

Author: shihab-dls

src/fastcs/controller.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections import Counter
-from collections.abc import Sequence
+from collections.abc import Iterator, Mapping, MutableMapping, Sequence
 from copy import deepcopy
 from typing import get_type_hints
 
@@ -17,6 +17,7 @@ class BaseController(Tracer):
 
     #: Attributes passed from the device at runtime.
     attributes: dict[str, Attribute]
+    root_attribute: Attribute | None = None
 
     description: str | None = None
 
@@ -36,7 +37,7 @@ def __init__(
         if not hasattr(self, "attributes"):
             self.attributes = {}
         self._path: list[str] = path or []
-        self.__sub_controller_tree: dict[str, Controller] = {}
+        self.__sub_controller_tree: dict[str, BaseController] = {}
 
         self._bind_attrs()
 
@@ -144,7 +145,7 @@ def add_attribute(self, name, attribute: Attribute):
         self.attributes[name] = attribute
         super().__setattr__(name, attribute)
 
-    def add_sub_controller(self, name: str, sub_controller: Controller):
+    def add_sub_controller(self, name: str, sub_controller: BaseController):
         if name in self.__sub_controller_tree.keys():
             raise ValueError(
                 f"Cannot add sub controller {sub_controller}. "
@@ -166,7 +167,7 @@ def add_sub_controller(self, name: str, sub_controller: Controller):
             self.attributes[name] = sub_controller.root_attribute
 
     @property
-    def sub_controllers(self) -> dict[str, Controller]:
+    def sub_controllers(self) -> dict[str, BaseController]:
         return self.__sub_controller_tree
 
     def __repr__(self):
@@ -194,17 +195,73 @@ class Controller(BaseController):
     such as generating a UI or creating parameters for a control system.
     """
 
-    root_attribute: Attribute | None = None
-
     def __init__(
         self,
         description: str | None = None,
         ios: Sequence[AttributeIO[T, AttributeIORefT]] | None = None,
     ) -> None:
         super().__init__(description=description, ios=ios)
 
+    def add_sub_controller(self, name: str, sub_controller: BaseController):
+        if name.isdigit():
+            raise ValueError(
+                f"Cannot add sub controller {name}. "
+                "Numeric-only names are not allowed; use ControllerVector instead"
+            )
+        return super().add_sub_controller(name, sub_controller)
+
     async def connect(self) -> None:
         pass
 
     async def disconnect(self) -> None:
         pass
+
+
+class ControllerVector(MutableMapping[int, Controller], BaseController):
+    """A controller with a collection of identical sub controllers distinguished
+    by a numeric value"""
+
+    def __init__(
+        self,
+        children: Mapping[int, Controller],
+        description: str | None = None,
+        ios: Sequence[AttributeIO[T, AttributeIORefT]] | None = None,
+    ) -> None:
+        super().__init__(description=description, ios=ios)
+        self._children: dict[int, Controller] = {}
+        for index, child in children.items():
+            self[index] = child
+
+    def add_sub_controller(self, name: str, sub_controller: BaseController):
+        raise NotImplementedError(
+            "Cannot add named sub controller to ControllerVector. "
+            "Use __setitem__ instead, for indexed sub controllers. "
+            "E.g., vector[1] = Controller()"
+        )
+
+    def __getitem__(self, key: int) -> Controller:
+        try:
+            return self._children[key]
+        except KeyError as exception:
+            raise KeyError(
+                f"ControllerVector does not have Controller with key {key}"
+            ) from exception
+
+    def __setitem__(self, key: int, value: Controller) -> None:
+        if not isinstance(key, int):
+            msg = f"Expected int, got {key}"
+            raise TypeError(msg)
+        if not isinstance(value, Controller):
+            msg = f"Expected Controller, got {value}"
+            raise TypeError(msg)
+        self._children[key] = value
+        super().add_sub_controller(str(key), value)
+
+    def __delitem__(self, key: int) -> None:
+        raise NotImplementedError("Cannot delete sub controller from ControllerVector.")
+
+    def __iter__(self) -> Iterator[int]:
+        yield from self._children
+
+    def __len__(self) -> int:
+        return len(self._children)

src/fastcs/controller_api.py

@@ -38,9 +38,12 @@ def walk_api(self) -> Iterator["ControllerAPI"]:
             yield from api.walk_api()
 
     def __repr__(self):
-        return f"""\
-ControllerAPI(path={self.path}, sub_apis=[{", ".join(self.sub_apis.keys())}])\
-"""
+        return (
+            f"ControllerAPI("
+            f"path={self.path}, "
+            f"sub_apis=[{', '.join(self.sub_apis.keys())}]"
+            f")"
+        )
 
     def get_scan_and_initial_coros(
         self,

src/fastcs/transport/epics/ca/ioc.py

@@ -113,9 +113,14 @@ def _add_sub_controller_pvi_info(pv_prefix: str, parent: ControllerAPI):
 
     for child in parent.sub_apis.values():
         child_pvi = f"{controller_pv_prefix(pv_prefix, child)}:PVI"
-        child_name = child.path[-1].lower()
+        child_name = (
+            f"__{child.path[-1]}"  # Sub-Controller of ControllerVector
+            if child.path[-1].isdigit()
+            else child.path[-1]
+        )
+
+        _add_pvi_info(child_pvi, parent_pvi, child_name.lower())
 
-        _add_pvi_info(child_pvi, parent_pvi, child_name)
         _add_sub_controller_pvi_info(pv_prefix, child)
 
 

src/fastcs/transport/epics/gui.py

@@ -160,6 +160,8 @@ def extract_api_components(self, controller_api: ControllerAPI) -> Tree:
         components: Tree = []
 
         for name, api in controller_api.sub_apis.items():
+            if name.isdigit():
+                name = f"{controller_api.path[-1]}{name}"
             components.append(
                 Group(
                     name=snake_to_pascal(name),

src/fastcs/transport/epics/pva/ioc.py

@@ -2,7 +2,7 @@
 
 from p4p.server import Server, StaticProvider
 
-from fastcs.attributes import Attribute, AttrR, AttrRW, AttrW
+from fastcs.attributes import AttrR, AttrRW, AttrW
 from fastcs.controller_api import ControllerAPI
 from fastcs.transport.epics.util import controller_pv_prefix
 from fastcs.util import snake_to_pascal
@@ -12,32 +12,23 @@
     make_shared_read_pv,
     make_shared_write_pv,
 )
-from .pvi_tree import AccessModeType, PviTree
-
-
-def _attribute_to_access(attribute: Attribute) -> AccessModeType:
-    match attribute:
-        case AttrRW():
-            return "rw"
-        case AttrR():
-            return "r"
-        case AttrW():
-            return "w"
-        case _:
-            raise ValueError(f"Unknown attribute type {type(attribute)}")
+from .pvi import add_pvi_info
 
 
 async def parse_attributes(
     root_pv_prefix: str, root_controller_api: ControllerAPI
-) -> list[StaticProvider]:
+) -> StaticProvider:
     """Parses `Attribute` s into p4p signals in handlers."""
-    pvi_tree = PviTree(root_pv_prefix)
     provider = StaticProvider(root_pv_prefix)
 
     for controller_api in root_controller_api.walk_api():
         pv_prefix = controller_pv_prefix(root_pv_prefix, controller_api)
-
-        pvi_tree.add_sub_device(pv_prefix, controller_api.description)
+        add_pvi_info(
+            provider=provider,
+            pv_prefix=pv_prefix,
+            controller_api=controller_api,
+            description=controller_api.description,
+        )
 
         for attr_name, attribute in controller_api.attributes.items():
             full_pv_name = f"{pv_prefix}:{snake_to_pascal(attr_name)}"
@@ -47,23 +38,19 @@ async def parse_attributes(
                     attribute_pv_rbv = make_shared_read_pv(attribute)
                     provider.add(f"{full_pv_name}", attribute_pv)
                     provider.add(f"{full_pv_name}_RBV", attribute_pv_rbv)
-                    pvi_tree.add_signal(f"{full_pv_name}", "rw")
                 case AttrR():
                     attribute_pv = make_shared_read_pv(attribute)
                     provider.add(f"{full_pv_name}", attribute_pv)
-                    pvi_tree.add_signal(f"{full_pv_name}", "r")
                 case AttrW():
                     attribute_pv = make_shared_write_pv(attribute)
                     provider.add(f"{full_pv_name}", attribute_pv)
-                    pvi_tree.add_signal(f"{full_pv_name}", "w")
 
         for attr_name, method in controller_api.command_methods.items():
             full_pv_name = f"{pv_prefix}:{snake_to_pascal(attr_name)}"
             command_pv = make_command_pv(method.fn)
             provider.add(f"{full_pv_name}", command_pv)
-            pvi_tree.add_signal(f"{full_pv_name}", "x")
 
-    return [provider, pvi_tree.make_provider()]
+    return provider
 
 
 class P4PIOC:
@@ -74,8 +61,8 @@ def __init__(self, pv_prefix: str, controller_api: ControllerAPI):
         self.controller_api = controller_api
 
     async def run(self):
-        providers = await parse_attributes(self.pv_prefix, self.controller_api)
+        provider = await parse_attributes(self.pv_prefix, self.controller_api)
 
         endless_event = asyncio.Event()
-        with Server(providers):
+        with Server([provider]):
             await endless_event.wait()

src/fastcs/transport/epics/pva/pvi.py

@@ -0,0 +1,118 @@
+from collections import defaultdict
+from typing import Literal
+
+from p4p import Type, Value
+from p4p.nt.common import alarm, timeStamp
+from p4p.server import StaticProvider
+from p4p.server.asyncio import SharedPV
+
+from fastcs.attributes import Attribute, AttrR, AttrRW, AttrW
+from fastcs.controller_api import ControllerAPI
+from fastcs.util import snake_to_pascal
+
+from .types import p4p_alarm_states, p4p_timestamp_now
+
+AccessModeType = Literal["r", "w", "rw", "d", "x"]
+
+
+# TODO: This should be removed after https://github.com/DiamondLightSource/FastCS/issues/260
+def _attribute_to_access(attribute: Attribute) -> AccessModeType:
+    match attribute:
+        case AttrRW():
+            return "rw"
+        case AttrR():
+            return "r"
+        case AttrW():
+            return "w"
+        case _:
+            raise ValueError(f"Unknown attribute type {type(attribute)}")
+
+
+def add_pvi_info(
+    provider: StaticProvider,
+    pv_prefix: str,
+    controller_api: ControllerAPI,
+    description: str | None = None,
+) -> None:
+    """Add PVI information to given provider."""
+    provider.add(
+        f"{pv_prefix}:PVI",
+        SharedPV(initial=_make_p4p_value(pv_prefix, controller_api, description)),
+    )
+
+
+def _make_p4p_value(
+    pv_prefix: str, controller_api: ControllerAPI, description: str | None
+) -> Value:
+    display = (
+        {"display": {"description": description}} if description is not None else {}
+    )  # Defined here so the value can be (none)
+
+    raw_value = _make_p4p_raw_value(pv_prefix, controller_api)
+    p4p_type = _make_type_for_raw_value(raw_value)
+
+    try:
+        return Value(
+            p4p_type,
+            {
+                **p4p_alarm_states(),
+                **p4p_timestamp_now(),
+                **display,
+                "value": raw_value,
+            },
+        )
+    except KeyError as e:
+        raise ValueError(f"Failed to create p4p Value from {raw_value}") from e
+
+
+def _make_p4p_raw_value(pv_prefix: str, controller_api: ControllerAPI) -> dict:
+    p4p_raw_value = defaultdict(dict)
+    # Sub-controller api returned if current item is a Controller
+    for pv_leaf, sub_controller_api in controller_api.sub_apis.items():
+        # Add Controller entry
+        pv = f"{pv_prefix}:{snake_to_pascal(pv_leaf)}:PVI"
+        if sub_controller_api.path[-1].isdigit():
+            # Sub-device of a ControllerVector
+            p4p_raw_value[f"__{int(pv_leaf)}"]["d"] = pv
+        else:
+            p4p_raw_value[pv_leaf]["d"] = pv
+    for pv_leaf, attribute in controller_api.attributes.items():
+        # Add attribute entry
+        pv = f"{pv_prefix}:{snake_to_pascal(pv_leaf)}"
+        p4p_raw_value[pv_leaf][_attribute_to_access(attribute)] = pv
+    for pv_leaf, _ in controller_api.command_methods.items():
+        pv = f"{pv_prefix}:{snake_to_pascal(pv_leaf)}"
+        p4p_raw_value[pv_leaf]["x"] = pv
+
+    return p4p_raw_value
+
+
+def _make_type_for_raw_value(raw_value: dict) -> Type:
+    p4p_raw_type = []
+    for pvi_group_name, access_to_field in raw_value.items():
+        pvi_group_structure = []
+        for access, field in access_to_field.items():
+            if isinstance(field, str):
+                pvi_group_structure.append((access, "s"))
+            elif isinstance(field, dict):
+                pvi_group_structure.append(
+                    (
+                        access,
+                        (
+                            "S",
+                            None,
+                            [(v, "s") for v, _ in field.items()],
+                        ),
+                    )
+                )
+
+        p4p_raw_type.append((pvi_group_name, ("S", "structure", pvi_group_structure)))
+
+    return Type(
+        [
+            ("alarm", alarm),
+            ("timeStamp", timeStamp),
+            ("display", ("S", None, [("description", "s")])),
+            ("value", ("S", "structure", p4p_raw_type)),
+        ]
+    )