Improve component graph's filtering

While we are changing the interface of component retrieval in the
component graph, we take the opportunity to make its interface more
generic.

We now accept any Iterable for `matching_ids` and `matching_types`, as
well as one individual items. Most uses of it filter by a single item
anyway so it is annoying to have to wrap it in a Iterable.

Most uses of those methods were updated to use the single-item overload.

Signed-off-by: Leandro Lucarella <[email protected]>

Author: llucax

benchmarks/power_distribution/power_distributor.py

@@ -143,7 +143,7 @@ async def run() -> None:
     )
 
     all_batteries = connection_manager.get().component_graph.components(
-        matching_types={Battery},
+        matching_types=Battery
     )
     batteries_ids = {c.id for c in all_batteries}
     # Take some time to get data from components

src/frequenz/sdk/microgrid/_old_component_data.py

@@ -160,7 +160,7 @@ def _check_category(cls, component_id: ComponentId) -> None:
         from .. import microgrid
 
         components = microgrid.connection_manager.get().component_graph.components(
-            matching_ids={component_id}
+            matching_ids=component_id
         )
         if not components:
             raise ValueError(f"Unable to find component with {component_id}")

src/frequenz/sdk/microgrid/_power_distributing/_component_managers/_battery_manager.py

@@ -146,7 +146,7 @@ def __init__(
         self._results_sender = results_sender
         self._api_power_request_timeout = api_power_request_timeout
         self._batteries = connection_manager.get().component_graph.components(
-            matching_types={Battery}
+            matching_types=Battery
         )
         self._battery_ids = {battery.id for battery in self._batteries}
 

src/frequenz/sdk/microgrid/_power_distributing/_component_managers/_ev_charger_manager/_ev_charger_manager.py

@@ -112,7 +112,7 @@ def _get_ev_charger_ids(self) -> collections.abc.Set[ComponentId]:
         return {
             evc.id
             for evc in connection_manager.get().component_graph.components(
-                matching_types={EvCharger}
+                matching_types=EvCharger
             )
         }
 

src/frequenz/sdk/microgrid/_power_distributing/_component_managers/_pv_inverter_manager/_pv_inverter_manager.py

@@ -255,6 +255,6 @@ def _get_pv_inverter_ids(self) -> collections.abc.Set[ComponentId]:
         return {
             inv.id
             for inv in connection_manager.get().component_graph.components(
-                matching_types={SolarInverter}
+                matching_types=SolarInverter
             )
         }

src/frequenz/sdk/microgrid/_power_wrapper.py

@@ -87,7 +87,7 @@ def _start_power_managing_actor(self) -> None:
         # Currently the power managing actor only supports batteries.  The below
         # constraint needs to be relaxed if the actor is extended to support other
         # components.
-        if not component_graph.components(matching_types={self._component_class}):
+        if not component_graph.components(matching_types=self._component_class):
             _logger.warning(
                 "No %s found in the component graph. "
                 "The power managing actor will not be started.",
@@ -119,7 +119,7 @@ def _start_power_distributing_actor(self) -> None:
             return
 
         component_graph = connection_manager.get().component_graph
-        if not component_graph.components(matching_types={self._component_class}):
+        if not component_graph.components(matching_types=self._component_class):
             _logger.warning(
                 "No %s found in the component graph. "
                 "The power distributing actor will not be started.",

src/frequenz/sdk/microgrid/component_graph.py

@@ -24,7 +24,7 @@
 import asyncio
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Set
 
 import networkx as nx
 from frequenz.client.common.microgrid.components import ComponentId
@@ -66,8 +66,8 @@ class ComponentGraph(ABC):
     @abstractmethod
     def components(
         self,
-        matching_ids: set[ComponentId] | None = None,
-        matching_types: set[type[Component]] | None = None,
+        matching_ids: Iterable[ComponentId] | ComponentId | None = None,
+        matching_types: Iterable[type[Component]] | type[Component] | None = None,
     ) -> set[Component]:
         """Fetch the components of the microgrid.
 
@@ -386,8 +386,8 @@ def __init__(
     @override
     def components(
         self,
-        matching_ids: set[ComponentId] | None = None,
-        matching_types: set[type[Component]] | None = None,
+        matching_ids: Iterable[ComponentId] | ComponentId | None = None,
+        matching_types: Iterable[type[Component]] | type[Component] | None = None,
     ) -> set[Component]:
         """Fetch the components of the microgrid.
 
@@ -399,6 +399,22 @@ def components(
             The set of components currently connected to the microgrid, filtered by
                 the provided `matching_ids` and `matching_types` values.
         """
+        match matching_ids:
+            case ComponentId():
+                matching_ids = {matching_ids}
+            case Set():
+                pass
+            case Iterable():
+                matching_ids = set(matching_ids)
+
+        match matching_types:
+            case type():
+                matching_types = {matching_types}
+            case Set():
+                pass
+            case Iterable():
+                matching_types = set(matching_types)
+
         selection: Iterable[Component]
         selection_ids = (
             self._graph.nodes

src/frequenz/sdk/timeseries/battery_pool/_battery_pool_reference_store.py

@@ -142,7 +142,7 @@ def _get_all_batteries(self) -> frozenset[ComponentId]:
         """
         graph = connection_manager.get().component_graph
         return frozenset(
-            battery.id for battery in graph.components(matching_types={Battery})
+            battery.id for battery in graph.components(matching_types=Battery)
         )
 
     async def _update_battery_status(

src/frequenz/sdk/timeseries/ev_charger_pool/_ev_charger_pool_reference_store.py

@@ -76,7 +76,7 @@ def __init__(  # pylint: disable=too-many-arguments
         else:
             graph = connection_manager.get().component_graph
             self.component_ids = frozenset(
-                {evc.id for evc in graph.components(matching_types={EvCharger})}
+                {evc.id for evc in graph.components(matching_types=EvCharger)}
             )
 
         self.power_bounds_subs: dict[str, asyncio.Task[None]] = {}

src/frequenz/sdk/timeseries/formula_engine/_formula_generators/_chp_power_formula.py

@@ -71,7 +71,7 @@ def _get_chp_meters(self) -> abc.Set[ComponentId]:
             FormulaGenerationError: If there's no dedicated meter attached to every CHP.
         """
         component_graph = connection_manager.get().component_graph
-        chps = component_graph.components(matching_types={Chp})
+        chps = component_graph.components(matching_types=Chp)
 
         chp_meters: set[ComponentId] = set()
         for chp in chps: