fix(graphs): keep edge_weights/node_weights/coords in sync after add_edges_from

qoolqit/graphs/base_graph.py

@@ -64,6 +64,40 @@ def _reset_dicts(self) -> None:
         self._node_weights = {n: None for n in self.nodes}
         self._edge_weights = {e: None for e in self.sorted_edges}
 
+    def add_edges_from(self, ebunch_to_add: Iterable, **attr: Any) -> None:
+        """Add all the edges in ebunch_to_add, keeping weight/coord dicts in sync.
+
+        `nx.Graph.add_edges_from` is used directly by the NetworkX layer and,
+        unlike the alternative constructors here, does not go through
+        `_reset_dicts()`. Left unhandled, edges (and any nodes they
+        implicitly create) added after construction are silently missing
+        from `_edge_weights`, `_node_weights` and `_coords` -- see #431.
+
+        Any newly-added node or edge is seeded with a weight/coordinate of
+        None, matching `_reset_dicts()`'s convention. Existing entries are
+        left untouched.
+
+        Arguments:
+            ebunch_to_add: container of edges, as accepted by NetworkX.
+            attr: edge data (or labels or objects) assigned via keyword
+                arguments, as accepted by NetworkX.
+        """
+        super().add_edges_from(ebunch_to_add, **attr)
+        if not hasattr(self, "_edge_weights"):
+            # Called from __init__ (via the initial `self.add_edges_from(edges)`),
+            # before `_reset_dicts()` has run for the first time. The constructor
+            # performs the initial sync itself right after, so there is nothing
+            # to do here yet.
+            return
+        for n in self.nodes:
+            if n not in self._node_weights:
+                self._node_weights[n] = None
+            if n not in self._coords:
+                self._coords[n] = None
+        for e in self.sorted_edges:
+            if e not in self._edge_weights:
+                self._edge_weights[e] = None
+
     @classmethod
     def from_nodes(cls, nodes: Iterable) -> BaseGraph:
         """Construct a base graph from a set of nodes.

tests/test_graphs/test_base_graph.py

@@ -210,9 +210,9 @@ def test_from_matrix(n_nodes: int) -> None:
 def test_to_matrix_unweighted(n_nodes: int) -> None:
     graph = BaseGraph.from_nodes(range(n_nodes))
     graph.add_edges_from(random_edge_list(range(n_nodes), k=2 * n_nodes))
-    # FIXME: _edge_weights is a snapshot that goes stale after add_edges_from;
-    # see issue #431 (edge_weights does not reflect edges added after construction).
-    graph._reset_dicts()
+    # add_edges_from() now keeps _node_weights/_edge_weights in sync on its
+    # own (see issue #431), so the _reset_dicts() workaround that used to be
+    # needed here is gone.
     assert not graph.has_node_weights
     assert not graph.has_edge_weights
 
@@ -231,6 +231,38 @@ def test_to_matrix_unweighted(n_nodes: int) -> None:
         assert matrix[j, i] == 0.0
 
 
+def test_add_edges_from_updates_weights_and_coords() -> None:
+    """Regression test for #431.
+
+    `edge_weights` (and `node_weights`/`coords`, for the same reason) must
+    reflect edges added via `add_edges_from` after construction, not just
+    the edges present at construction time.
+    """
+    graph = BaseGraph.from_nodes([0, 1, 2])
+
+    # Edges added post-construction must show up with a None weight, not be
+    # silently missing.
+    graph.add_edges_from([(0, 1), (1, 2)])
+    assert graph.edge_weights == {(0, 1): None, (1, 2): None}
+    assert graph.node_weights == {0: None, 1: None, 2: None}
+
+    # Existing weights must not be clobbered by a later add_edges_from call.
+    graph.edge_weights = {(0, 1): 1.5, (1, 2): 2.5}
+    graph.node_weights = {0: 0.1, 1: 0.2, 2: 0.3}
+    graph.add_edges_from([(0, 2)])
+    assert graph.edge_weights == {(0, 1): 1.5, (1, 2): 2.5, (0, 2): None}
+    assert graph.node_weights == {0: 0.1, 1: 0.2, 2: 0.3}
+
+    # Edges that implicitly introduce brand-new nodes must extend
+    # node_weights/coords for those nodes too, instead of leaving them
+    # missing (which would otherwise surface later as a KeyError).
+    graph._coords = {0: (0.0, 0.0), 1: (1.0, 0.0), 2: (2.0, 0.0)}
+    graph.add_edges_from([(2, 3)])
+    assert 3 in graph.node_weights and graph.node_weights[3] is None
+    assert 3 in graph.coords and graph.coords[3] is None
+    assert graph.edge_weights[(2, 3)] is None
+
+
 @pytest.mark.parametrize("n_nodes", [5, 10, 50])
 def test_to_matrix_weighted(n_nodes: int) -> None:
     graph = BaseGraph.from_nodes(range(n_nodes))