Store the original Component and Connection in the component graph

For some reason the Component and Connection objects were converted to
`dict` to store them in the component graph. This requires re-building
the objects when they need to be extracted, which is wasteful.

It also means that type information is lost, and in the future, when
adopting the microgrid API client using v0.17, we have a class hierarchy
for components instead of relaying in the category and type, so we would
even need to store the type if we wanted to rebuild the objects.

Having nested objects is also an issue, as then the nested properties
were kept as `dict`, adding an inconsistency and type unsafety.

With this change we are able to get rid of a hack to retrieve the grid
metadata.

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

Author: llucax

src/frequenz/sdk/microgrid/component_graph.py

@@ -22,11 +22,9 @@
 """
 
 import asyncio
-import dataclasses
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable
-from dataclasses import asdict
 
 import networkx as nx
 from frequenz.client.microgrid import (
@@ -42,6 +40,10 @@
 # pylint: disable=too-many-lines
 
 
+# Constant to store the actual obejcts as data attached to the graph nodes and edges
+_DATA_KEY = "data"
+
+
 class InvalidGraphError(Exception):
     """Exception type that will be thrown if graph data is not valid."""
 
@@ -398,18 +400,16 @@ def components(
             Set of the components currently connected to the microgrid, filtered by
                 the provided `component_ids` and `component_categories` values.
         """
-        if component_ids is None:
-            # If any node has not node[1], then it will not pass validations step.
-            selection: Iterable[Component] = map(
-                lambda node: Component(**(node[1])), self._graph.nodes(data=True)
-            )
-        else:
-            valid_ids = filter(self._graph.has_node, component_ids)
-            selection = map(lambda idx: Component(**self._graph.nodes[idx]), valid_ids)
+        selection: Iterable[Component]
+        selection_ids = (
+            self._graph.nodes
+            if component_ids is None
+            else component_ids & self._graph.nodes
+        )
+        selection = (self._graph.nodes[i][_DATA_KEY] for i in selection_ids)
 
         if component_categories is not None:
-            types: set[ComponentCategory] = component_categories
-            selection = filter(lambda c: c.category in types, selection)
+            selection = filter(lambda c: c.category in component_categories, selection)
 
         return set(selection)
 
@@ -430,19 +430,19 @@ def connections(
             Set of the connections between components in the microgrid, filtered by
                 the provided `start`/`end` choices.
         """
-        if start is None:
-            if end is None:
-                selection = self._graph.edges
-            else:
-                selection = self._graph.in_edges(end)
-
-        else:
-            selection = self._graph.out_edges(start)
-            if end is not None:
-                end_ids: set[int] = end
-                selection = filter(lambda c: c[1] in end_ids, selection)
-
-        return set(map(lambda c: Connection(c[0], c[1]), selection))
+        match (start, end):
+            case (None, None):
+                selection_ids = self._graph.edges
+            case (None, _):
+                selection_ids = self._graph.in_edges(end)
+            case (_, None):
+                selection_ids = self._graph.out_edges(start)
+            case (_, _):
+                start_edges = self._graph.out_edges(start)
+                end_edges = self._graph.in_edges(end)
+                selection_ids = set(start_edges).intersection(end_edges)
+
+        return set(self._graph.edges[i][_DATA_KEY] for i in selection_ids)
 
     def predecessors(self, component_id: int) -> set[Component]:
         """Fetch the graph predecessors of the specified component.
@@ -466,9 +466,7 @@ def predecessors(self, component_id: int) -> set[Component]:
 
         predecessors_ids = self._graph.predecessors(component_id)
 
-        return set(
-            map(lambda idx: Component(**self._graph.nodes[idx]), predecessors_ids)
-        )
+        return set(map(lambda idx: self._graph.nodes[idx][_DATA_KEY], predecessors_ids))
 
     def successors(self, component_id: int) -> set[Component]:
         """Fetch the graph successors of the specified component.
@@ -492,7 +490,7 @@ def successors(self, component_id: int) -> set[Component]:
 
         successors_ids = self._graph.successors(component_id)
 
-        return set(map(lambda idx: Component(**self._graph.nodes[idx]), successors_ids))
+        return set(map(lambda idx: self._graph.nodes[idx][_DATA_KEY], successors_ids))
 
     def refresh_from(
         self,
@@ -526,9 +524,14 @@ def refresh_from(
 
         new_graph = nx.DiGraph()
         for component in components:
-            new_graph.add_node(component.component_id, **asdict(component))
+            new_graph.add_node(component.component_id, **{_DATA_KEY: component})
 
-        new_graph.add_edges_from(dataclasses.astuple(c) for c in connections)
+        # Store the original connection object in the edge data (third item in the
+        # tuple) so that we can retrieve it later.
+        for connection in connections:
+            new_graph.add_edge(
+                connection.start, connection.end, **{_DATA_KEY: connection}
+            )
 
         # check if we can construct a valid ComponentGraph
         # from the new NetworkX graph data
@@ -908,8 +911,9 @@ def _validate_graph(self) -> None:
         if not nx.is_directed_acyclic_graph(self._graph):
             raise InvalidGraphError("Component graph is not a tree!")
 
-        # node[0] is required by the graph definition
-        # If any node has not node[1], then it will not pass validations step.
+        # This check doesn't seem to have much sense, it only search for nodes without
+        # data associated with them. We leave it here for now, but we should consider
+        # removing it in the future.
         undefined: list[int] = [
             node[0] for node in self._graph.nodes(data=True) if len(node[1]) == 0
         ]

src/frequenz/sdk/timeseries/grid.py

@@ -207,17 +207,12 @@ def initialize(
         if grid_connections[0].metadata is None:
             raise RuntimeError("Grid metadata is None")
 
-        # The current implementation of the Component Graph fails to
-        # effectively convert components from a dictionary representation to
-        # the expected Component object.
-        # Specifically for the component metadata, it hands back a dictionary
-        # instead of the expected ComponentMetadata type.
-        metadata = grid_connections[0].metadata
-        if isinstance(metadata, dict):
-            if fuse_dict := metadata.get("fuse", None):
-                fuse = Fuse(
-                    max_current=Current.from_amperes(fuse_dict.get("max_current", 0.0))
+        if grid_connections[0].metadata.fuse is not None:
+            fuse = Fuse(
+                max_current=Current.from_amperes(
+                    grid_connections[0].metadata.fuse.max_current
                 )
+            )
 
         if fuse is None:
             _logger.warning("The grid connection point does not have a fuse")