Issue #22: Store associations between components

docs/explanation/system.md

@@ -58,3 +58,41 @@ Component.model_json_schema()
 - `infrasys` includes some basic quantities in [infrasys.quantities](#quantity-api).
 - Pint will automatically convert a list or list of lists of values into a `numpy.ndarray`.
 infrasys will handle serialization/de-serialization of these types.
+
+
+### Component Associations
+The system tracks associations between components in order to optimize lookups.
+
+For example, suppose a Generator class has a field for a Bus. It is trivial to find a generator's
+bus. However, if you need to find all generators connected to specific bus, you would have to
+traverse all generators in the system and check their bus values.
+
+Every time you add a component to a system, `infrasys` inspects the component type for composed
+components. It checks for directly connected components, such as `Generator.bus`, and lists of
+components. (It does not inspect other composite data structures like dictionaries.)
+
+`infrasys` stores these component associations in a SQLite table and so lookups are fast.
+
+Here is how to complete this example:
+
+```python
+generators = system.list_parent_components(bus)
+```
+
+If you only want to find specific types, you can pass that type as well.
+```python
+generators = system.list_parent_components(bus, component_type=Generator)
+```
+
+**Warning**: There is one potentially problematic case.
+
+Suppose that you have a system with generators and buses and then reassign the buses, as in
+```
+gen1.bus = other_bus
+```
+
+`infrasys` cannot detect such reassignments and so the component associations will be incorrect.
+You must inform `infrasys` to rebuild its internal table.
+```
+system.rebuild_component_associations()
+```

src/infrasys/component_associations.py

@@ -0,0 +1,99 @@
+import sqlite3
+from typing import Optional, Type
+from uuid import UUID
+
+from loguru import logger
+
+from infrasys.component import Component
+from infrasys.utils.sqlite import execute
+
+
+class ComponentAssociations:
+    """Stores associations between components. Allows callers to quickly find components composed
+    by other components, such as the generator to which a bus is connected."""
+
+    TABLE_NAME = "component_associations"
+
+    def __init__(self) -> None:
+        # This uses a different database because it is not persisted when the system
+        # is saved to files. It will be rebuilt during de-serialization.
+        self._con = sqlite3.connect(":memory:")
+        self._create_metadata_table()
+
+    def _create_metadata_table(self):
+        schema = [
+            "id INTEGER PRIMARY KEY",
+            "component_uuid TEXT",
+            "component_type TEXT",
+            "attached_component_uuid TEXT",
+            "attached_component_type TEXT",
+        ]
+        schema_text = ",".join(schema)
+        cur = self._con.cursor()
+        execute(cur, f"CREATE TABLE {self.TABLE_NAME}({schema_text})")
+        execute(cur, f"CREATE INDEX by_ac_uuid ON {self.TABLE_NAME}(attached_component_uuid)")
+        self._con.commit()
+        logger.debug("Created in-memory component associations table")
+
+    def add(self, *components: Component):
+        """Store an association between each component and directly attached subcomponents.
+
+        - Inspects the type of each field of each component's type. Looks for subtypes of
+          Component and lists of subtypes of Component.
+        - Does not consider component fields that are dictionaries or other data structures.
+        """
+        rows = []
+        for component in components:
+            for field in type(component).model_fields:
+                val = getattr(component, field)
+                if isinstance(val, Component):
+                    rows.append(self._make_row(component, val))
+                elif isinstance(val, list) and val and isinstance(val[0], Component):
+                    for item in val:
+                        rows.append(self._make_row(component, item))
+                # FUTURE: consider supporting dictionaries like these examples:
+                # dict[str, Component]
+                # dict[str, [Component]]
+
+        if rows:
+            self._insert_rows(rows)
+
+    def clear(self) -> None:
+        """Clear all component associations."""
+        execute(self._con.cursor(), f"DELETE FROM {self.TABLE_NAME}")
+        logger.info("Cleared all component associations.")
+
+    def list_parent_components(
+        self, component: Component, component_type: Optional[Type[Component]] = None
+    ) -> list[UUID]:
+        """Return a list of all component UUIDS that compose this component.
+        For example, return all components connected to a bus.
+        """
+        where_clause = "WHERE attached_component_uuid = ?"
+        if component_type is None:
+            params = [str(component.uuid)]
+        else:
+            params = [str(component.uuid), component_type.__name__]
+            where_clause += " AND component_type = ?"
+        query = f"SELECT component_uuid FROM {self.TABLE_NAME} {where_clause}"
+        cur = self._con.cursor()
+        return [UUID(x[0]) for x in execute(cur, query, params)]
+
+    def _insert_rows(self, rows: list[tuple]) -> None:
+        cur = self._con.cursor()
+        placeholder = ",".join(["?"] * len(rows[0]))
+        query = f"INSERT INTO {self.TABLE_NAME} VALUES({placeholder})"
+        try:
+            cur.executemany(query, rows)
+        finally:
+            self._con.commit()
+
+    @staticmethod
+    def _make_row(component: Component, attached_component: Component):
+        return (
+            None,
+            str(component.uuid),
+            type(component).__name__,
+            str(attached_component.uuid),
+            type(attached_component).__name__,
+        )

src/infrasys/component_manager.py

@@ -1,13 +1,19 @@
 """Manages components"""
 
-from collections import defaultdict
 import itertools
-from typing import Any, Callable, Iterable, Type
+from collections import defaultdict
+from typing import Any, Callable, Iterable, Optional, Type
 from uuid import UUID
 from loguru import logger
 
 from infrasys.component import Component
-from infrasys.exceptions import ISAlreadyAttached, ISNotStored, ISOperationNotAllowed
+from infrasys.component_associations import ComponentAssociations
+from infrasys.exceptions import (
+    ISAlreadyAttached,
+    ISNotStored,
+    ISOperationNotAllowed,
+    ISInvalidParameter,
+)
 from infrasys.models import make_label, get_class_and_name_from_label
 
 
@@ -23,6 +29,7 @@ def __init__(
         self._components_by_uuid: dict[UUID, Component] = {}
         self._uuid = uuid
         self._auto_add_composed_components = auto_add_composed_components
+        self._associations = ComponentAssociations()
 
     @property
     def auto_add_composed_components(self) -> bool:
@@ -34,17 +41,23 @@ def auto_add_composed_components(self, val: bool) -> None:
         """Set auto_add_composed_components."""
         self._auto_add_composed_components = val
 
-    def add(self, *args: Component, deserialization_in_progress=False) -> None:
+    def add(self, *components: Component, deserialization_in_progress=False) -> None:
         """Add one or more components to the system.
 
         Raises
         ------
         ISAlreadyAttached
             Raised if a component is already attached to a system.
         """
-        for component in args:
+        if not components:
+            msg = "add_associations requires at least one component"
+            raise ISInvalidParameter(msg)
+
+        for component in components:
             self._add(component, deserialization_in_progress)
 
+        self._associations.add(*components)
+
     def get(self, component_type: Type[Component], name: str) -> Any:
         """Return the component with the passed type and name.
 
@@ -167,8 +180,22 @@ def iter_all(self) -> Iterable[Any]:
         """Return an iterator over all components."""
         return self._components_by_uuid.values()
 
+    def list_parent_components(
+        self, component: Component, component_type: Optional[Type[Component]] = None
+    ) -> list[Component]:
+        """Return a list of all components that compose this component."""
+        return [
+            self.get_by_uuid(x)
+            for x in self._associations.list_parent_components(
+                component, component_type=component_type
+            )
+        ]
+
     def to_records(
-        self, component_type: Type[Component], filter_func: Callable | None = None, **kwargs
+        self,
+        component_type: Type[Component],
+        filter_func: Callable | None = None,
+        **kwargs,
     ) -> Iterable[dict]:
         """Return a dictionary representation of the requested components.
 
@@ -207,6 +234,15 @@ def remove(self, component: Component) -> Any:
             msg = f"{component.label} is not stored"
             raise ISNotStored(msg)
 
+        attached_components = self.list_parent_components(component)
+        if attached_components:
+            label = ", ".join((x.label for x in attached_components))
+            msg = (
+                f"Cannot remove {component.label} because it is attached to these components: "
+                f"{label}"
+            )
+            raise ISOperationNotAllowed(msg)
+
         container = self._components[component_type][component.name]
         for i, comp in enumerate(container):
             if comp.uuid == component.uuid:
@@ -259,6 +295,14 @@ def change_uuid(self, component: Component) -> None:
         msg = "change_component_uuid"
         raise NotImplementedError(msg)
 
+    def rebuild_component_associations(self) -> None:
+        """Clear the component associations and rebuild the table. This may be necessary
+        if a user reassigns connected components that are part of a system.
+        """
+        self._associations.clear()
+        self._associations.add(*self.iter_all())
+        logger.info("Rebuilt all component associations.")
+
     def update(
         self,
         component_type: Type[Component],
@@ -292,6 +336,7 @@ def _add(self, component: Component, deserialization_in_progress: bool) -> None:
 
         self._components[cls][name].append(component)
         self._components_by_uuid[component.uuid] = component
+
         logger.debug("Added {} to the system", component.label)
 
     def _check_component_addition(self, component: Component) -> None:
@@ -303,7 +348,7 @@ def _check_component_addition(self, component: Component) -> None:
                 self._handle_composed_component(val)
                 # Recurse.
                 self._check_component_addition(val)
-            if isinstance(val, list) and val and isinstance(val[0], Component):
+            elif isinstance(val, list) and val and isinstance(val[0], Component):
                 for item in val:
                     self._handle_composed_component(item)
                     # Recurse.

src/infrasys/exceptions.py

@@ -18,13 +18,17 @@ class ISFileExists(ISBaseException):
 
 
 class ISConflictingArguments(ISBaseException):
-    """Raised if the arguments are conflict."""
+    """Raised if the arguments conflict."""
 
 
 class ISConflictingSystem(ISBaseException):
     """Raised if the system has conflicting values."""
 
 
+class ISInvalidParameter(ISBaseException):
+    """Raised if a parameter is invalid."""
+
+
 class ISNotStored(ISBaseException):
     """Raised if the requested object is not stored."""
 

src/infrasys/system.py

@@ -597,6 +597,22 @@ def get_component_types(self) -> Iterable[Type[Component]]:
         """
         return self._component_mgr.get_types()
 
+    def list_parent_components(
+        self, component: Component, component_type: Optional[Type[Component]] = None
+    ) -> list[Component]:
+        """Return a list of all components that compose this component.
+
+        An example usage is where you need to find all components connected to a bus and the Bus
+        class does not contain that information. The system tracks these connections internally
+        and can find those components quickly.
+
+        Examples
+        --------
+        >>> components = system.list_parent_components(bus)
+        >>> print(f"These components are connected to {bus.label}: ", " ".join(components))
+        """
+        return self._component_mgr.list_parent_components(component, component_type=component_type)
+
     def list_components_by_name(self, component_type: Type[Component], name: str) -> list[Any]:
         """Return all components that match component_type and name.
 
@@ -625,6 +641,12 @@ def iter_all_components(self) -> Iterable[Any]:
         """
         return self._component_mgr.iter_all()
 
+    def rebuild_component_associations(self) -> None:
+        """Clear the component associations and rebuild the table. This may be necessary
+        if a user reassigns connected components that are part of a system.
+        """
+        self._component_mgr.rebuild_component_associations()
+
     def remove_component(self, component: Component) -> Any:
         """Remove the component from the system and return it.
 
@@ -636,6 +658,8 @@ def remove_component(self, component: Component) -> Any:
         ------
         ISNotStored
             Raised if the component is not stored in the system.
+        ISOperationNotAllowed
+            Raised if the other components hold references to this component.
 
         Examples
         --------