Repository: NREL/infrasys

Issue:
Add an argument to `list_child_components` and `list_parent_components` that accepts and AbstractClass

Discussion:

Solution patch:
diff --git a/src/infrasys/component_associations.py b/src/infrasys/component_associations.py
index ac45e35..a2eb53d 100644
--- a/src/infrasys/component_associations.py
+++ b/src/infrasys/component_associations.py
@@ -5,6 +5,7 @@
 from loguru import logger
 
 from infrasys.component import Component
+from infrasys.utils.classes import get_all_concrete_subclasses
 from infrasys.utils.sqlite import execute
 
 
@@ -74,11 +75,11 @@ def list_child_components(
         For example, return the bus attached to a generator.
         """
         where_clause = "WHERE component_uuid = ?"
-        if component_type is None:
-            params = [str(component.uuid)]
-        else:
-            params = [str(component.uuid), component_type.__name__]
-            where_clause += " AND attached_component_type = ?"
+        params = [str(component.uuid)]
+        if component_type is not None:
+            res = _make_params_and_where_clause(component_type, "attached_component_type")
+            params.extend(res[0])
+            where_clause += res[1]
         query = f"SELECT attached_component_uuid FROM {self.TABLE_NAME} {where_clause}"
         cur = self._con.cursor()
         return [UUID(x[0]) for x in execute(cur, query, params)]
@@ -90,11 +91,11 @@ def list_parent_components(
         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 = ?"
+        params = [str(component.uuid)]
+        if component_type is not None:
+            res = _make_params_and_where_clause(component_type, "component_type")
+            params.extend(res[0])
+            where_clause += res[1]
         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)]
@@ -128,3 +129,22 @@ def _make_row(component: Component, attached_component: Component):
             str(attached_component.uuid),
             type(attached_component).__name__,
         )
+
+
+def _make_params_and_where_clause(
+    component_type: Type[Component], field: str
+) -> tuple[list[str], str]:
+    params = _make_type_params(component_type)
+    if len(params) == 1:
+        where_clause = f" AND {field} = ?"
+    else:
+        where_clause = f" AND {field} IN ({','.join(['?'] * len(params))})"
+    return params, where_clause
+
+
+def _make_type_params(component_type: Type[Component]) -> list[str]:
+    params: list[str] = []
+    subclasses = get_all_concrete_subclasses(component_type) or [component_type]
+    for cls in subclasses:
+        params.append(cls.__name__)  # type: ignore
+    return params
diff --git a/src/infrasys/system.py b/src/infrasys/system.py
index 7b7c6aa..6f95d2d 100644
--- a/src/infrasys/system.py
+++ b/src/infrasys/system.py
@@ -744,6 +744,7 @@ def list_child_components(
         component: Component
         component_type: Optional[Type[Component]]
             Filter the returned list to components of this type.
+            If the type has subclasses, all subclasses will be included.
 
         See Also
         --------
@@ -765,6 +766,7 @@ class does not contain that information. The system tracks these connections int
         component: Component
         component_type: Optional[Type[Component]]
             Filter the returned list to components of this type.
+            If the type has subclasses, all subclasses will be included.
 
         Examples
         --------
diff --git a/src/infrasys/utils/classes.py b/src/infrasys/utils/classes.py
new file mode 100644
index 0000000..7f5234b
--- /dev/null
+++ b/src/infrasys/utils/classes.py
@@ -0,0 +1,28 @@
+"""Utility functions for classes."""
+
+from typing import Type
+
+
+_cached_classes: dict[Type, set[Type]] = {}
+
+
+def get_all_concrete_subclasses(cls: Type) -> set[Type]:
+    """Return all concrete subclasses of a class recursively.
+    This excludes subclasses that have their own subclasses.
+    Note that infrasys component types should not be used as both abstract and concrete, even
+    though Python allows it.
+    This keeps a cache and so dynamic changes to the class hierarchy will not be reflected.
+    """
+    cached_cls = _cached_classes.get(cls)
+    if cached_cls is not None:
+        return cached_cls
+
+    subclasses = set()
+    for subclass in cls.__subclasses__():
+        if subclass.__subclasses__():
+            subclasses.update(get_all_concrete_subclasses(subclass))
+        else:
+            subclasses.add(subclass)
+
+    _cached_classes[cls] = subclasses
+    return subclasses


Test patch:
diff --git a/tests/test_system.py b/tests/test_system.py
index cfb051f..a319378 100644
--- a/tests/test_system.py
+++ b/tests/test_system.py
@@ -157,12 +157,12 @@ def test_component_associations(tmp_path):
         subsystem = SimpleSubsystem(name=f"test-subsystem{i}", generators=[gen1, gen2])
         system.add_components(geo, bus, gen1, gen2, subsystem)
 
-    def check_attached_components(my_sys):
+    def check_attached_components(my_sys, parent_type, child_type):
         for i in range(3):
             bus = my_sys.get_component(SimpleBus, f"bus{i}")
             gen1 = my_sys.get_component(SimpleGenerator, f"gen{i}a")
             gen2 = my_sys.get_component(SimpleGenerator, f"gen{i}b")
-            attached = my_sys.list_parent_components(bus, component_type=SimpleGenerator)
+            attached = my_sys.list_parent_components(bus, component_type=parent_type)
             assert len(attached) == 2
             labels = {gen1.label, gen2.label}
             for component in attached:
@@ -172,24 +172,28 @@ def check_attached_components(my_sys):
                 assert attached_subsystems[0].name == f"test-subsystem{i}"
                 assert not my_sys.list_parent_components(attached_subsystems[0])
                 assert my_sys.list_child_components(component) == [bus]
-                assert my_sys.list_child_components(component, component_type=SimpleBus) == [bus]
+                assert my_sys.list_child_components(component, component_type=child_type) == [bus]
 
             for component in (bus, gen1, gen2):
                 with pytest.raises(ISOperationNotAllowed):
                     my_sys.remove_component(component)
 
-    check_attached_components(system)
+    check_attached_components(system, SimpleGenerator, SimpleBus)
+    check_attached_components(system, GeneratorBase, Component)
+    check_attached_components(system, Component, Component)
     system._component_mgr._associations.clear()
     for component in system.iter_all_components():
         assert not system.list_parent_components(component)
 
     system.rebuild_component_associations()
-    check_attached_components(system)
+    check_attached_components(system, SimpleGenerator, SimpleBus)
+    check_attached_components(system, GeneratorBase, Component)
 
     save_dir = tmp_path / "test_system"
     system.save(save_dir)
     system2 = SimpleSystem.from_json(save_dir / "system.json")
-    check_attached_components(system2)
+    check_attached_components(system2, SimpleGenerator, SimpleBus)
+    check_attached_components(system2, GeneratorBase, Component)
 
     bus = system2.get_component(SimpleBus, "bus1")
     with pytest.raises(ISOperationNotAllowed):
diff --git a/tests/test_utils_classes.py b/tests/test_utils_classes.py
new file mode 100644
index 0000000..0d0873b
--- /dev/null
+++ b/tests/test_utils_classes.py
@@ -0,0 +1,20 @@
+from infrasys import Component
+from infrasys.utils.classes import get_all_concrete_subclasses
+
+
+class Level1(Component):
+    field1: int
+
+
+class Level2(Level1):
+    field2: int
+
+
+class Level3(Level2):
+    field3: int
+
+
+def test_get_all_concrete_subclasses():
+    assert get_all_concrete_subclasses(Level1) == {Level3}
+    assert get_all_concrete_subclasses(Level2) == {Level3}
+    assert not get_all_concrete_subclasses(Level3)