Interface updates

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

pyproject.toml

@@ -73,6 +73,8 @@ dev-pylint = [
   "frequenz-microgrid-component-graph[dev-mkdocs,dev-noxfile,dev-pytest]",
 ]
 dev-pytest = [
+  "frequenz-microgrid-component-graph[microgrid]",
+  "frequenz-microgrid-component-graph[assets]",
   "pytest == 8.4.2",
   "pylint == 4.0.2",       # We need this to check for the examples
   "frequenz-repo-config[extra-lint-examples] == 0.13.6",

python/frequenz/microgrid_component_graph/__init__.pyi

@@ -6,7 +6,7 @@
 """A graph representation of the electrical components in a microgrid."""
 
 from collections import abc
-from typing import AbstractSet, Generic, Protocol, TypeVar
+from typing import Generic, Protocol, TypeVar
 
 class InvalidGraphError(Exception):
     """Exception type that will be thrown if graph data is not valid."""
@@ -60,14 +60,15 @@ class ComponentGraphConfig:
 
 ComponentT = TypeVar("ComponentT", bound=ComponentProtocol)
 ConnectionT = TypeVar("ConnectionT", bound=ConnectionProtocol)
+ComponentIdT = TypeVar("ComponentIdT", bound=ComponentIdProtocol)
 
-class ComponentGraph(Generic[ComponentT, ConnectionT]):
+class ComponentGraph(Generic[ComponentT, ConnectionT, ComponentIdT]):
     """A graph representation of the electrical components in a microgrid."""
 
     def __init__(
         self,
-        components: AbstractSet[ComponentT],
-        connections: AbstractSet[ConnectionT],
+        components: abc.Iterable[ComponentT],
+        connections: abc.Iterable[ConnectionT],
         config: ComponentGraphConfig = ComponentGraphConfig(),
     ) -> None:
         """Initialize this instance.
@@ -78,7 +79,7 @@ class ComponentGraph(Generic[ComponentT, ConnectionT]):
             config: The configuration for the component graph.
         """
 
-    def component(self, component_id: ComponentIdProtocol) -> ComponentT:
+    def component(self, component_id: ComponentIdT) -> ComponentT:
         """Fetch the component with the specified `component_id`.
 
         Args:
@@ -93,23 +94,23 @@ class ComponentGraph(Generic[ComponentT, ConnectionT]):
 
     def components(
         self,
-        filter_by_ids: set[ComponentIdProtocol] | None = None,
-        filter_by_types: set[type[ComponentProtocol]] | None = None,
+        matching_ids: abc.Iterable[ComponentIdT] | None = None,
+        matching_types: abc.Iterable[type[ComponentT]] | None = None,
     ) -> set[ComponentT]:
         """Fetch all components in this graph.
 
         Returns:
             A set of all components in this graph.
         """
 
-    def connections(self) -> set[ConnectionT]:
+    def connections(self) -> abc.Set[ConnectionT]:
         """Fetch all connections in this graph.
 
         Returns:
             A set of all connections in this graph.
         """
 
-    def predecessors(self, component_id: ComponentIdProtocol) -> set[ComponentT]:
+    def predecessors(self, component_id: ComponentIdT) -> abc.Set[ComponentT]:
         """Fetch all predecessors of the specified component ID.
 
         Args:
@@ -122,7 +123,7 @@ class ComponentGraph(Generic[ComponentT, ConnectionT]):
             ValueError: if no component exists with the given ID.
         """
 
-    def successors(self, component_id: ComponentIdProtocol) -> set[ComponentT]:
+    def successors(self, component_id: ComponentIdT) -> abc.Set[ComponentT]:
         """Fetch all successors of the specified component ID.
 
         Args:
@@ -135,7 +136,7 @@ class ComponentGraph(Generic[ComponentT, ConnectionT]):
             ValueError: if no component exists with the given ID.
         """
 
-    def is_pv_meter(self, component_id: ComponentIdProtocol) -> bool:
+    def is_pv_meter(self, component_id: ComponentIdT) -> bool:
         """Check if the specified component is a PV meter.
 
         Args:
@@ -148,7 +149,7 @@ class ComponentGraph(Generic[ComponentT, ConnectionT]):
             ValueError: if no component exists with the given ID.
         """
 
-    def is_battery_meter(self, component_id: ComponentIdProtocol) -> bool:
+    def is_battery_meter(self, component_id: ComponentIdT) -> bool:
         """Check if the specified component is a battery meter.
 
         Args:
@@ -161,7 +162,7 @@ class ComponentGraph(Generic[ComponentT, ConnectionT]):
             ValueError: if no component exists with the given ID.
         """
 
-    def is_ev_charger_meter(self, component_id: ComponentIdProtocol) -> bool:
+    def is_ev_charger_meter(self, component_id: ComponentIdT) -> bool:
         """Check if the specified component is an EV charger meter.
 
         Args:
@@ -174,7 +175,7 @@ class ComponentGraph(Generic[ComponentT, ConnectionT]):
             ValueError: if no component exists with the given ID.
         """
 
-    def is_chp_meter(self, component_id: ComponentIdProtocol) -> bool:
+    def is_chp_meter(self, component_id: ComponentIdT) -> bool:
         """Check if the specified component is a CHP meter.
 
         Args:
@@ -208,30 +209,28 @@ class ComponentGraph(Generic[ComponentT, ConnectionT]):
             The grid formula as a string.
         """
 
-    def pv_formula(self, pv_inverter_ids: abc.Set[ComponentIdProtocol] | None) -> str:
+    def pv_formula(self, pv_inverter_ids: abc.Set[ComponentIdT] | None) -> str:
         """Generate the PV formula for this component graph.
 
         Returns:
             The PV formula as a string.
         """
 
-    def battery_formula(self, battery_ids: abc.Set[ComponentIdProtocol] | None) -> str:
+    def battery_formula(self, battery_ids: abc.Set[ComponentIdT] | None) -> str:
         """Generate the battery formula for this component graph.
 
         Returns:
             The battery formula as a string.
         """
 
-    def chp_formula(self, chp_ids: abc.Set[ComponentIdProtocol] | None) -> str:
+    def chp_formula(self, chp_ids: abc.Set[ComponentIdT] | None) -> str:
         """Generate the CHP formula for this component graph.
 
         Returns:
             The CHP formula as a string.
         """
 
-    def ev_charger_formula(
-        self, ev_charger_ids: abc.Set[ComponentIdProtocol] | None
-    ) -> str:
+    def ev_charger_formula(self, ev_charger_ids: abc.Set[ComponentIdT] | None) -> str:
         """Generate the EV charger formula for this component graph.
 
         Returns:
@@ -246,17 +245,15 @@ class ComponentGraph(Generic[ComponentT, ConnectionT]):
         """
 
     def battery_coalesce_formula(
-        self, battery_ids: abc.Set[ComponentIdProtocol] | None
+        self, battery_ids: abc.Set[ComponentIdT] | None
     ) -> str:
         """Generate the battery coalesce formula for this component graph.
 
         Returns:
             The battery coalesced formula as a string.
         """
 
-    def pv_coalesce_formula(
-        self, pv_inverter_ids: abc.Set[ComponentIdProtocol] | None
-    ) -> str:
+    def pv_coalesce_formula(self, pv_inverter_ids: abc.Set[ComponentIdT] | None) -> str:
         """Generate the PV coalesce formula for this component graph.
 
         Returns:

src/graph.rs

@@ -106,16 +106,16 @@ impl ComponentGraph {
         cls.into()
     }
 
-    #[pyo3(signature = (filter_by_ids=None, filter_by_types=None))]
+    #[pyo3(signature = (matching_ids=None, matching_types=None))]
     fn components(
         &self,
-        filter_by_ids: Option<Bound<'_, PyAny>>,
-        filter_by_types: Option<Bound<'_, PyAny>>,
+        matching_ids: Option<Bound<'_, PyAny>>,
+        matching_types: Option<Bound<'_, PyAny>>,
     ) -> PyResult<Py<PySet>> {
         let iter = self.graph.components();
 
         Python::attach(|py| {
-            let components: Vec<_> = if let Some(component_categories) = filter_by_types {
+            let components: Vec<_> = if let Some(component_categories) = matching_types {
                 let categories: Vec<cg::ComponentCategory> = component_categories
                     .try_iter()?
                     .map(|c| category_from_python_component(py, &c?))
@@ -127,7 +127,7 @@ impl ComponentGraph {
                 iter.collect()
             };
 
-            let components: Vec<_> = if let Some(ids) = filter_by_ids {
+            let components: Vec<_> = if let Some(ids) = matching_ids {
                 let ids_set: BTreeSet<u64> = ids
                     .try_iter()?
                     .map(|id| extract_int::<u64>(py, id?))

tests/test_microgrid_component_graph.py

@@ -6,6 +6,7 @@
 from frequenz.client.common.microgrid import MicrogridId
 from frequenz.client.common.microgrid.components import ComponentId
 from frequenz.client.microgrid.component import (
+    Component,
     ComponentConnection,
     GridConnectionPoint,
     Meter,
@@ -16,10 +17,14 @@
 
 def test_loading() -> None:
     """Test that the microgrid_component_graph module loads correctly."""
-    graph = microgrid_component_graph.ComponentGraph(
+    graph: microgrid_component_graph.ComponentGraph[
+        Component, ComponentConnection, ComponentId
+    ] = microgrid_component_graph.ComponentGraph(
         components={
             GridConnectionPoint(
-                id=ComponentId(1), microgrid_id=MicrogridId(1), rated_fuse_current=100
+                id=ComponentId(1),
+                microgrid_id=MicrogridId(1),
+                rated_fuse_current=100,
             ),
             Meter(id=ComponentId(2), microgrid_id=MicrogridId(1)),
         },