Allow passing single elements to component graph connections()

Like with component graph `components()`, now we support passing a
single element to `matching_sources` and `matching_destinations`.

The `components()` function is also simplified to make filtering more
compact.

Also updates uses with a single element to remove the unnecessary
container.

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

Author: llucax

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, Set
+from collections.abc import Callable, Iterable
 
 import networkx as nx
 from frequenz.client.common.microgrid.components import ComponentId
@@ -83,8 +83,8 @@ def components(
     @abstractmethod
     def connections(
         self,
-        matching_sources: set[ComponentId] | None = None,
-        matching_destinations: set[ComponentId] | None = None,
+        matching_sources: Iterable[ComponentId] | ComponentId | None = None,
+        matching_destinations: Iterable[ComponentId] | ComponentId | None = None,
     ) -> set[ComponentConnection]:
         """Fetch the connections between microgrid components.
 
@@ -399,27 +399,15 @@ 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)
+        matching_ids = _comp_ids_to_iter(matching_ids)
+        if isinstance(matching_types, type):
+            matching_types = {matching_types}
 
         selection: Iterable[Component]
         selection_ids = (
             self._graph.nodes
             if matching_ids is None
-            else matching_ids & self._graph.nodes
+            else set(matching_ids) & self._graph.nodes
         )
         selection = (self._graph.nodes[i][_DATA_KEY] for i in selection_ids)
 
@@ -433,8 +421,8 @@ def components(
     @override
     def connections(
         self,
-        matching_sources: set[ComponentId] | None = None,
-        matching_destinations: set[ComponentId] | None = None,
+        matching_sources: Iterable[ComponentId] | ComponentId | None = None,
+        matching_destinations: Iterable[ComponentId] | ComponentId | None = None,
     ) -> set[ComponentConnection]:
         """Fetch the connections between microgrid components.
 
@@ -447,6 +435,9 @@ def connections(
             The set of connections between components in the microgrid, filtered by
                 the provided `matching_sources` and `matching_destinations` choices.
         """
+        matching_sources = _comp_ids_to_iter(matching_sources)
+        matching_destinations = _comp_ids_to_iter(matching_destinations)
+
         match (matching_sources, matching_destinations):
             case (None, None):
                 selection = self._graph.edges
@@ -1128,3 +1119,11 @@ def _validate_leaf_components(self) -> None:
             raise InvalidGraphError(
                 f"Leaf components with graph successors: {with_successors}"
             )
+
+
+def _comp_ids_to_iter(
+    ids: Iterable[ComponentId] | ComponentId | None,
+) -> Iterable[ComponentId] | None:
+    if isinstance(ids, ComponentId):
+        return (ids,)
+    return ids

tests/microgrid/test_graph.py

@@ -389,18 +389,18 @@ def test_connection_filters(self) -> None:
 
         # with start filter applied, we get back only connections whose `start`
         # component matches one of the provided IDs
-        assert graph.connections(matching_sources={ComponentId(8)}) == set()
-        assert graph.connections(matching_sources={ComponentId(7)}) == set()
-        assert graph.connections(matching_sources={ComponentId(6)}) == set()
-        assert graph.connections(matching_sources={ComponentId(5)}) == set()
-        assert graph.connections(matching_sources={ComponentId(4)}) == set()
-        assert graph.connections(matching_sources={ComponentId(3)}) == set()
-        assert graph.connections(matching_sources={ComponentId(2)}) == {
+        assert graph.connections(matching_sources=ComponentId(8)) == set()
+        assert graph.connections(matching_sources=ComponentId(7)) == set()
+        assert graph.connections(matching_sources=ComponentId(6)) == set()
+        assert graph.connections(matching_sources=ComponentId(5)) == set()
+        assert graph.connections(matching_sources=ComponentId(4)) == set()
+        assert graph.connections(matching_sources=ComponentId(3)) == set()
+        assert graph.connections(matching_sources=ComponentId(2)) == {
             ComponentConnection(source=ComponentId(2), destination=ComponentId(4)),
             ComponentConnection(source=ComponentId(2), destination=ComponentId(5)),
             ComponentConnection(source=ComponentId(2), destination=ComponentId(6)),
         }
-        assert graph.connections(matching_sources={ComponentId(1)}) == {
+        assert graph.connections(matching_sources=ComponentId(1)) == {
             ComponentConnection(source=ComponentId(1), destination=ComponentId(2)),
             ComponentConnection(source=ComponentId(1), destination=ComponentId(3)),
         }
@@ -427,23 +427,23 @@ def test_connection_filters(self) -> None:
 
         # with end filter applied, we get back only connections whose `end`
         # component matches one of the provided IDs
-        assert graph.connections(matching_destinations={ComponentId(8)}) == set()
-        assert graph.connections(matching_destinations={ComponentId(6)}) == {
+        assert graph.connections(matching_destinations=ComponentId(8)) == set()
+        assert graph.connections(matching_destinations=ComponentId(6)) == {
             ComponentConnection(source=ComponentId(2), destination=ComponentId(6))
         }
-        assert graph.connections(matching_destinations={ComponentId(5)}) == {
+        assert graph.connections(matching_destinations=ComponentId(5)) == {
             ComponentConnection(source=ComponentId(2), destination=ComponentId(5))
         }
-        assert graph.connections(matching_destinations={ComponentId(4)}) == {
+        assert graph.connections(matching_destinations=ComponentId(4)) == {
             ComponentConnection(source=ComponentId(2), destination=ComponentId(4))
         }
-        assert graph.connections(matching_destinations={ComponentId(3)}) == {
+        assert graph.connections(matching_destinations=ComponentId(3)) == {
             ComponentConnection(source=ComponentId(1), destination=ComponentId(3))
         }
-        assert graph.connections(matching_destinations={ComponentId(2)}) == {
+        assert graph.connections(matching_destinations=ComponentId(2)) == {
             ComponentConnection(source=ComponentId(1), destination=ComponentId(2))
         }
-        assert graph.connections(matching_destinations={ComponentId(1)}) == set()
+        assert graph.connections(matching_destinations=ComponentId(1)) == set()
         assert graph.connections(
             matching_destinations={ComponentId(1), ComponentId(2), ComponentId(3)}
         ) == {
@@ -470,18 +470,18 @@ def test_connection_filters(self) -> None:
             ComponentConnection(source=ComponentId(2), destination=ComponentId(4)),
             ComponentConnection(source=ComponentId(2), destination=ComponentId(6)),
         }
-        assert graph.connections(matching_destinations={ComponentId(1)}) == set()
+        assert graph.connections(matching_destinations=ComponentId(1)) == set()
 
         # when both filters are applied, they are combined via AND logic, i.e.
         # a connection must have its `start` matching one of the provided start
         # values, and its `end` matching one of the provided end values
         assert graph.connections(
-            matching_sources={ComponentId(1)}, matching_destinations={ComponentId(2)}
+            matching_sources=ComponentId(1), matching_destinations=ComponentId(2)
         ) == {ComponentConnection(source=ComponentId(1), destination=ComponentId(2))}
         assert (
             graph.connections(
-                matching_sources={ComponentId(2)},
-                matching_destinations={ComponentId(3)},
+                matching_sources=ComponentId(2),
+                matching_destinations=ComponentId(3),
             )
             == set()
         )