refactor: Refactor builder registry

src/techui_builder/builder.py

@@ -1,7 +1,6 @@
 import logging
 import os
 import re
-from collections import defaultdict
 from dataclasses import dataclass, field
 from pathlib import Path
 
@@ -34,9 +33,7 @@ class Builder:
 
     techui: Path = field(default=Path("techui.yaml"))
 
-    entities: defaultdict[str, list[Entity]] = field(
-        default_factory=lambda: defaultdict(list), init=False
-    )
+    entities: dict[str, Entity] = field(default_factory=dict, init=False)
     _services_dir: Path = field(init=False, repr=False)
     _write_directory: Path = field(init=False, repr=False)
 
@@ -180,10 +177,21 @@ def _extract_entities(self, service_name: str, service_yaml: Path):
                                 macros=macros,
                             )
 
-                            pv_root = prefix.split(":", maxsplit=1)[0]
-                            self.entities[pv_root].append(new_entity)
+                            self.entities[prefix] = new_entity
                     break
 
+    def get_entities_for_prefix(self, prefix: str) -> list[Entity]:
+        """Return entities matching a full PV prefix, or all entities under a root."""
+        results = []
+        if prefix in self.entities:
+            # Exact match (no suffix) — also catches root entity in loop below
+            results.append(self.entities[prefix])
+        # Children with suffixes (e.g. BL19I-MO-DCM-01:MOT1)
+        for entity in self.entities.values():
+            if entity.base_prefix == prefix and entity is not self.entities.get(prefix):
+                results.append(entity)
+        return results
+
     def _generate_screen(self, screen_name: str):
         self.generator.build_screen(screen_name)
         self.generator.write_screen(screen_name, self._write_directory)
@@ -210,41 +218,39 @@ def create_screens(self):
         for component_name, component in self.conf.components.items():
             screen_entities: list[Entity] = []
 
-            # ONLY IF there is a matching component and entity, generate a screen
-            if component.prefix in self.entities.keys():
-                # Populate child labels for any entities
-                # with the same prefix as the component
-                for entity in self.entities[component.prefix]:
-                    entity.child_labels = component.child_labels
-
-                screen_entities.extend(self.entities[component.prefix])
-
-                if component.extras is not None:
-                    # If component has any extras, add them to the entries to generate
-                    for extra_p in component.extras:
-                        if extra_p not in self.entities.keys():
-                            logger_.error(
-                                f"Extra prefix {extra_p} for {component_name} does not"
-                                " exist."
-                            )
-                            continue
-                        screen_entities.extend(self.entities[extra_p])
-
-                # This is used by both generate and validate,
-                # so called beforehand for tidiness
-                self.generator.build_widgets(component_name, screen_entities)
-                self.generator.build_groups(component_name, self.conf.components)
-
-                screens_to_validate = list(self.validator.validate.keys())
-
-                if component_name in screens_to_validate:
-                    self._validate_screen(component_name)
-                else:
-                    self._generate_screen(component_name)
-
-            else:
+            matched = self.get_entities_for_prefix(component.prefix)
+            if not matched:
                 logger_.warning(
-                    f"{self.techui.name}: The prefix [bold]{component.prefix}[/bold] "
-                    f"set in the component [bold]{component_name}[/bold] does not match"
-                    " any P field in the ioc.yaml files in services"
+                    f"{self.techui.name}: prefix [bold]{component.prefix}[/bold] "
+                    f"for component [bold]{component_name}[/bold] matched no entities."
                 )
+                continue
+
+            # Populate child labels for any entities
+            for entity in matched:
+                entity.child_labels = component.child_labels
+
+            screen_entities.extend(matched)
+
+            if component.extras is not None:
+                for extra_p in component.extras:
+                    extra_matched = self.get_entities_for_prefix(extra_p)
+                    if not extra_matched:
+                        logger_.error(
+                            f"Extra prefix {extra_p} for {component_name} matched no"
+                            " entities."
+                        )
+                        continue
+                    screen_entities.extend(extra_matched)
+
+            # This is used by both generate and validate,
+            # so called beforehand for tidiness
+            self.generator.build_widgets(component_name, screen_entities)
+            self.generator.build_groups(component_name, self.conf.components)
+
+            screens_to_validate = list(self.validator.validate.keys())
+
+            if component_name in screens_to_validate:
+                self._validate_screen(component_name)
+            else:
+                self._generate_screen(component_name)

src/techui_builder/generate.py

@@ -2,7 +2,6 @@
 import os
 import re
 from collections import defaultdict
-from collections.abc import Mapping
 from dataclasses import dataclass, field
 from pathlib import Path
 
@@ -12,7 +11,7 @@
 from phoebusgen import widget as pwidget
 from phoebusgen.widget.widgets import ActionButton, EmbeddedDisplay, Group
 
-from techui_builder.models import Component, Entity, TechUiSupport
+from techui_builder.models import Component, Entity, SupportEntityScreen, TechUiSupport
 
 logger_ = logging.getLogger(__name__)
 
@@ -159,11 +158,14 @@ def _update_macros(self, component: Entity) -> tuple[str, dict[str, str]]:
                 :2
             ]
             component_name = suffix.removeprefix(":").removesuffix(":")
-            suffix_key = next(k for k, v in component.macros.items() if v == suffix)
+            suffix_key = next(
+                (k for k, v in component.macros.items() if v == suffix), None
+            )
         except (IndexError, ValueError):
             prefix = component.prefix
             component_name = component.type
-            suffix_key = suffix = ""
+            suffix_key = None
+            suffix = ""
 
         # Try to get name from child labels if they exist,
         # if not, just use the name as it is.
@@ -175,27 +177,27 @@ def _update_macros(self, component: Entity) -> tuple[str, dict[str, str]]:
         prefix_key = next(k for k, v in component.macros.items() if v == prefix)
 
         new_macros[prefix_key] = prefix
-        if suffix_key != "":
+        if suffix_key is not None:
             new_macros[suffix_key] = suffix
             new_macros["label"] = component_name
 
         return component_name, new_macros
 
     def _allocate_widget(
-        self, screen_mapping: Mapping, component: Entity
+        self, screen_mapping: SupportEntityScreen, component: Entity
     ) -> EmbeddedDisplay | ActionButton | None | list[EmbeddedDisplay | ActionButton]:
         component_name, updated_macros = self._update_macros(component)
         component.macros.update(updated_macros)
 
         # Get relative path to screen
-        file = Template(screen_mapping["file"]).render(component.macros)
+        file = Template(screen_mapping.file).render(component.macros)
         if file.startswith("$(IOC)"):
             screen_path = support_screen_path = file.replace(
                 "$(IOC)", f"{self.beamline_url}/{component.service_name}"
             )  # Only works with related displays as
             # embedded displays need to access the file to get dimensions
 
-            assert screen_mapping["type"] == "related", (
+            assert screen_mapping.type == "related", (
                 "Only related displays can have remote screens"
             )
         else:
@@ -207,27 +209,34 @@ def _allocate_widget(
                 self.synoptic_dir, walk_up=True
             )
 
-        # For Gui Components with multiple components embedded, we add a suffix field
-        # to the components, and adjust the name and suffix accordingly
+        # For Gui Components with multiple components embedded, we add macro overrides
+        # via screen_macros and adjust the name/label accordingly
         try:
-            if screen_mapping["suffixes"] is not None:
-                suffix_dict: dict[str, str] = screen_mapping["suffixes"]
-                for suffix_key, suffix in suffix_dict.items():
-                    component.macros[suffix_key] = suffix
+            if screen_mapping.screen_macros:
+                for macro_key, macro_val in screen_mapping.screen_macros.items():
+                    component.macros[macro_key] = macro_val
 
                 # If no child label was specified...
-                if self.label_flag is False:
-                    # TODO: think of a better fallback component name for this
-                    component_name = (
-                        list(suffix_dict.values())[0]
-                        .removeprefix(":")
-                        .removesuffix(":")
-                    )
-                    component.macros["label"] = component_name
+                if not self.label_flag:
+                    # Prefer an explicit "label" or "suffix" key in screen_macros,
+                    # otherwise fall back to the first value (by insertion order).
+                    fallback_keys = ("label", "suffix")
+                    first_val = None
+                    for key in fallback_keys:
+                        if key in screen_mapping.screen_macros:
+                            first_val = screen_mapping.screen_macros[key]
+                            break
+                    if first_val is None:
+                        first_val = next(
+                            iter(screen_mapping.screen_macros.values()), None
+                        )
+                    if first_val is not None:
+                        component_name = first_val.removeprefix(":").removesuffix(":")
+                        component.macros["label"] = component_name
         except KeyError:
             pass
 
-        if screen_mapping["type"] == "embedded":
+        if screen_mapping.type == "embedded":
             height, width = self._get_screen_dimensions(str(screen_path))
             new_widget = pwidget.EmbeddedDisplay(
                 component_name,

src/techui_builder/models.py

@@ -307,6 +307,30 @@ class Entity(BaseModel):
         Field(description="Macros for the matching screen (can be empty)"),
     ]
 
+    @computed_field(repr=False, return_type=str)
+    @property
+    def base_prefix(self) -> str:
+        """PV prefix, excluding  any colons or suffix, e.g. 'BL01T-MO-MOTOR-01'."""
+        return self.prefix.split(":", maxsplit=1)[0]
+
+
+class SupportEntityScreen(BaseModel):
+    """One screen entry for a support module."""
+
+    file: Annotated[str, Field(description="Path to the .bob or .pvi.bob screen file")]
+    type: Annotated[str, Field(description="Screen type, e.g. 'embedded' or 'related'")]
+    screen_macros: Annotated[
+        dict[str, str] | None,
+        Field(
+            description="Screen-specific macro overrides",
+            # Make sure vscode is aware of schema validation
+            json_schema_extra={
+                "type": "object",
+                "additionalProperties": {"type": "string"},
+            },
+        ),
+    ] = None
+
 
 class SupportEntity(BaseModel):
     """
@@ -319,8 +343,8 @@ class SupportEntity(BaseModel):
         Field(description="Macros for the matching screen (can be empty)"),
     ]
     screens: Annotated[
-        list[dict[str, str | dict[str, str]]],
-        Field(description="Dictionary of available screens for the support module"),
+        list[SupportEntityScreen],
+        Field(description="List of available screens for the support module"),
     ]
 
 

src/techui_builder/utils.py

@@ -6,6 +6,15 @@
 
 logger_ = logging.getLogger(__name__)
 
+__all__ = [
+    "WidgetType",
+    "get_widgets",
+    "read_bob",
+    "_get_action_group",
+    "_get_macros",
+    "_get_nav_tabs",
+]
+
 
 class WidgetType(StrEnum):
     """Widget types in a .bob file."""