Fix pickling of graphs with contracted nodes (#1505)

* set node_removed if nodes contracted

* format

Author: d-miketa

src/digraph.rs

@@ -3007,6 +3007,7 @@ impl PyDiGraph {
             }
             (None, true) => self.graph.contract_nodes(nodes, obj, check_cycle)?,
         };
+        self.node_removed = true;
         Ok(res.index())
     }
 

src/graph.rs

@@ -1899,6 +1899,7 @@ impl PyGraph {
             }
             (None, true) => self.graph.contract_nodes(nodes, obj),
         };
+        self.node_removed = true;
         Ok(res.index())
     }
 

tests/digraph/test_pickle.py

@@ -39,3 +39,26 @@ def test_weight_graph(self):
         self.assertEqual([1, 2, 3], gprime.node_indices())
         self.assertEqual(["B", "C", "D"], gprime.nodes())
         self.assertEqual({1: (1, 2, "B -> C"), 3: (3, 1, "D -> B")}, dict(gprime.edge_index_map()))
+
+    def test_contracted_nodes_pickle(self):
+        """Test pickle/unpickle of directed graphs with contracted nodes (issue #1503)"""
+        g = rx.PyDiGraph()
+        g.add_node("A")  # Node 0
+        g.add_node("B")  # Node 1
+        g.add_node("C")  # Node 2
+
+        # Contract nodes 0 and 1 into a new node
+        contracted_idx = g.contract_nodes([0, 1], "AB")
+        g.add_edge(2, contracted_idx, "C -> AB")
+
+        # Verify initial state
+        self.assertEqual([2, contracted_idx], g.node_indices())
+        self.assertEqual([(2, contracted_idx)], g.edge_list())
+
+        # Test pickle/unpickle
+        gprime = pickle.loads(pickle.dumps(g))
+
+        # Verify the unpickled graph matches
+        self.assertEqual(g.node_indices(), gprime.node_indices())
+        self.assertEqual(g.edge_list(), gprime.edge_list())
+        self.assertEqual(g.nodes(), gprime.nodes())

tests/graph/test_pickle.py

@@ -39,3 +39,46 @@ def test_weight_graph(self):
         self.assertEqual([1, 2, 3], gprime.node_indices())
         self.assertEqual(["B", "C", "D"], gprime.nodes())
         self.assertEqual({1: (1, 2, "B -> C"), 3: (3, 1, "D -> B")}, dict(gprime.edge_index_map()))
+
+    def test_contracted_nodes_pickle(self):
+        """Test pickle/unpickle of graphs with contracted nodes (issue #1503)"""
+        g = rx.PyGraph()
+        g.add_node("A")  # Node 0
+        g.add_node("B")  # Node 1
+        g.add_node("C")  # Node 2
+
+        # Contract nodes 0 and 1 into a new node
+        contracted_idx = g.contract_nodes([0, 1], "AB")
+        g.add_edge(2, contracted_idx, "C -> AB")
+
+        # Verify initial state
+        self.assertEqual([2, contracted_idx], g.node_indices())
+        self.assertEqual([(2, contracted_idx)], g.edge_list())
+
+        # Test pickle/unpickle
+        gprime = pickle.loads(pickle.dumps(g))
+
+        # Verify the unpickled graph matches
+        self.assertEqual(g.node_indices(), gprime.node_indices())
+        self.assertEqual(g.edge_list(), gprime.edge_list())
+        self.assertEqual(g.nodes(), gprime.nodes())
+
+    def test_contracted_nodes_with_weights_pickle(self):
+        """Test pickle/unpickle of graphs with contracted nodes and edge weights"""
+        g = rx.PyGraph()
+        g.add_nodes_from(["Node0", "Node1", "Node2", "Node3"])
+        g.add_edges_from([(0, 2, "edge_0_2"), (1, 3, "edge_1_3")])
+
+        # Contract multiple nodes
+        contracted_idx = g.contract_nodes([0, 1], "Contracted_0_1")
+        g.add_edge(contracted_idx, 2, "contracted_to_2")
+        g.add_edge(3, contracted_idx, "3_to_contracted")
+
+        # Test pickle/unpickle
+        gprime = pickle.loads(pickle.dumps(g))
+
+        # Verify complete graph state is preserved
+        self.assertEqual(g.node_indices(), gprime.node_indices())
+        self.assertEqual(g.edge_list(), gprime.edge_list())
+        self.assertEqual(g.nodes(), gprime.nodes())
+        self.assertEqual(dict(g.edge_index_map()), dict(gprime.edge_index_map()))