Merge pull request #91 from DiamondLightSource/89-change-how-we-searc…

…h-for-attributes

Change how we search for attributes

src/fastcs/controller.py

@@ -3,6 +3,7 @@
 from collections.abc import Iterator
 from copy import copy
 from dataclasses import dataclass
+from typing import get_type_hints
 
 from .attributes import Attribute
 from .cs_methods import Command, Put, Scan
@@ -19,9 +20,14 @@ class SingleMapping:
 
 
 class BaseController:
+    #: Attributes passed from the device at runtime.
+    attributes: dict[str, Attribute]
+
     def __init__(self, path: list[str] | None = None) -> None:
+        if not hasattr(self, "attributes"):
+            self.attributes = {}
         self._path: list[str] = path or []
-        self.__sub_controller_tree: dict[str, BaseController] = {}
+        self.__sub_controller_tree: dict[str, SubController] = {}
 
         self._bind_attrs()
 
@@ -37,11 +43,27 @@ def set_path(self, path: list[str]):
         self._path = path
 
     def _bind_attrs(self) -> None:
-        for attr_name in dir(self):
-            attr = getattr(self, attr_name)
+        # Using a dictionary instead of a set to maintain order.
+        class_dir = {key: None for key in dir(type(self))}
+        class_type_hints = get_type_hints(type(self))
+
+        for attr_name in {**class_dir, **class_type_hints}:
+            if attr_name == "root_attribute":
+                continue
+
+            attr = getattr(self, attr_name, None)
             if isinstance(attr, Attribute):
+                if (
+                    attr_name in self.attributes
+                    and self.attributes[attr_name] is not attr
+                ):
+                    raise ValueError(
+                        f"`{type(self).__name__}` has conflicting attribute "
+                        f"`{attr_name}` already present in the attributes dict."
+                    )
                 new_attribute = copy(attr)
                 setattr(self, attr_name, new_attribute)
+                self.attributes[attr_name] = new_attribute
 
     def register_sub_controller(self, name: str, sub_controller: SubController):
         if name in self.__sub_controller_tree.keys():
@@ -52,7 +74,16 @@ def register_sub_controller(self, name: str, sub_controller: SubController):
         self.__sub_controller_tree[name] = sub_controller
         sub_controller.set_path(self.path + [name])
 
-    def get_sub_controllers(self) -> dict[str, BaseController]:
+        if isinstance(sub_controller.root_attribute, Attribute):
+            if name in self.attributes:
+                raise TypeError(
+                    f"Cannot set SubController `{name}` root attribute "
+                    f"on the parent controller `{type(self).__name__}` "
+                    f"as it already has an attribute of that name."
+                )
+            self.attributes[name] = sub_controller.root_attribute
+
+    def get_sub_controllers(self) -> dict[str, SubController]:
         return self.__sub_controller_tree
 
     def get_controller_mappings(self) -> list[SingleMapping]:
@@ -69,7 +100,6 @@ def _get_single_mapping(controller: BaseController) -> SingleMapping:
     scan_methods: dict[str, Scan] = {}
     put_methods: dict[str, Put] = {}
     command_methods: dict[str, Command] = {}
-    attributes: dict[str, Attribute] = {}
     for attr_name in dir(controller):
         attr = getattr(controller, attr_name)
         match attr:
@@ -79,11 +109,15 @@ def _get_single_mapping(controller: BaseController) -> SingleMapping:
                 scan_methods[attr_name] = scan_method
             case WrappedMethod(fastcs_method=Command(enabled=True) as command_method):
                 command_methods[attr_name] = command_method
-            case Attribute(enabled=True):
-                attributes[attr_name] = attr
+
+    enabled_attributes = {
+        name: attribute
+        for name, attribute in controller.attributes.items()
+        if attribute.enabled
+    }
 
     return SingleMapping(
-        controller, scan_methods, put_methods, command_methods, attributes
+        controller, scan_methods, put_methods, command_methods, enabled_attributes
     )
 
 
@@ -113,5 +147,7 @@ class SubController(BaseController):
     it as part of a larger device.
     """
 
+    root_attribute: Attribute | None = None
+
     def __init__(self) -> None:
         super().__init__()

tests/test_controller.py

@@ -1,11 +1,13 @@
 import pytest
 
+from fastcs.attributes import AttrR
 from fastcs.controller import (
     Controller,
     SubController,
     _get_single_mapping,
     _walk_mappings,
 )
+from fastcs.datatypes import Int
 
 
 def test_controller_nesting():
@@ -33,3 +35,92 @@ def test_controller_nesting():
         ValueError, match=r"SubController is already registered under .*"
     ):
         controller.register_sub_controller("c", sub_controller)
+
+
+class SomeSubController(SubController):
+    def __init__(self):
+        super().__init__()
+
+    sub_attribute = AttrR(Int())
+
+    root_attribute = AttrR(Int())
+
+
+class SomeController(Controller):
+    annotated_attr: AttrR
+    annotated_attr_not_defined_in_init: AttrR[int]
+    equal_attr = AttrR(Int())
+    annotated_and_equal_attr: AttrR[int] = AttrR(Int())
+
+    def __init__(self, sub_controller: SubController):
+        self.attributes = {}
+
+        self.annotated_attr = AttrR(Int())
+        self.attr_on_object = AttrR(Int())
+
+        self.attributes["_attributes_attr"] = AttrR(Int())
+        self.attributes["_attributes_attr_equal"] = self.equal_attr
+
+        super().__init__()
+        self.register_sub_controller("sub_controller", sub_controller)
+
+
+def test_attribute_parsing():
+    sub_controller = SomeSubController()
+    controller = SomeController(sub_controller)
+
+    mapping_walk = _walk_mappings(controller)
+
+    controller_mapping = next(mapping_walk)
+    assert set(controller_mapping.attributes.keys()) == {
+        "_attributes_attr",
+        "annotated_attr",
+        "_attributes_attr_equal",
+        "annotated_and_equal_attr",
+        "equal_attr",
+        "sub_controller",
+    }
+
+    assert SomeController.equal_attr is not controller.equal_attr
+    assert (
+        SomeController.annotated_and_equal_attr
+        is not controller.annotated_and_equal_attr
+    )
+
+    sub_controller_mapping = next(mapping_walk)
+    assert sub_controller_mapping.attributes == {
+        "sub_attribute": sub_controller.sub_attribute,
+    }
+
+
+def test_attribute_in_both_class_and_get_attributes():
+    class FailingController(Controller):
+        duplicate_attribute = AttrR(Int())
+
+        def __init__(self):
+            self.attributes = {"duplicate_attribute": AttrR(Int())}
+            super().__init__()
+
+    with pytest.raises(
+        ValueError,
+        match=(
+            "`FailingController` has conflicting attribute `duplicate_attribute` "
+            "already present in the attributes dict."
+        ),
+    ):
+        FailingController()
+
+
+def test_root_attribute():
+    class FailingController(SomeController):
+        sub_controller = AttrR(Int())
+
+    with pytest.raises(
+        TypeError,
+        match=(
+            "Cannot set SubController `sub_controller` root attribute "
+            "on the parent controller `FailingController` as it already "
+            "has an attribute of that name."
+        ),
+    ):
+        next(_walk_mappings(FailingController(SomeSubController())))