Add APIs to update an edge's data payload inplace (#236)

* Add APIs to update an edge's data payload inplace

This commit adds 2 new methods to the graph classes, update_edge and
update_edge_by_index, which are used to update the data payload/weight
of an edge in the graph in place. Prior to this the only way to update
an edge would be to remove it and add a new edge with the new data
payload which may change the edge index (assuming something was using
it).

* Apply suggestions from code review

Co-authored-by: Lauren Capelluto <laurencapelluto@gmail.com>

* Add tests with parallel edges

Co-authored-by: Lauren Capelluto <laurencapelluto@gmail.com>

releasenotes/notes/update-edge-407745423983c801.yaml

@@ -0,0 +1,10 @@
+---
+features:
+  - |
+    Two new methods, :meth:`~retworkx.PyDiGraph.update_edge` and
+    :meth:`~retworkx.PyDiGraph.update_edge_by_index` were added to the
+    :class:`retworkx.PyDiGraph` and :class:`retworkx.PyGraph` (
+    :meth:`~retworkx.PyGraph.update_edge` and
+    :meth:`~retworkx.PyGraph.update_edge_by_index`) classes. These methods
+    are used to update the data payload/weight of an edge in the graph either
+    by the nodes of an edge or by edge index.

src/digraph.rs

@@ -611,6 +611,60 @@ impl PyDiGraph {
         Ok(data)
     }
 
+    /// Update an edge's weight/payload inplace
+    ///
+    /// If there are parallel edges in the graph only one edge will be updated.
+    /// if you need to update a specific edge or need to ensure all parallel
+    /// edges get updated you should use
+    /// :meth:`~retworkx.PyDiGraph.update_edge_by_index` instead.
+    ///
+    /// :param int source: The index for the first node
+    /// :param int target: The index for the second node
+    ///
+    /// :raises NoEdgeBetweenNodes: When there is no edge between nodes
+    #[text_signature = "(self, source, target, edge /)"]
+    pub fn update_edge(
+        &mut self,
+        source: usize,
+        target: usize,
+        edge: PyObject,
+    ) -> PyResult<()> {
+        let index_a = NodeIndex::new(source);
+        let index_b = NodeIndex::new(target);
+        let edge_index = match self.graph.find_edge(index_a, index_b) {
+            Some(edge_index) => edge_index,
+            None => {
+                return Err(NoEdgeBetweenNodes::new_err(
+                    "No edge found between nodes",
+                ))
+            }
+        };
+        let data = self.graph.edge_weight_mut(edge_index).unwrap();
+        *data = edge;
+        Ok(())
+    }
+
+    /// Update an edge's weight/payload by the edge index
+    ///
+    /// :param int edge_index: The index for the edge
+    /// :param object edge: The data payload/weight to update the edge with
+    ///
+    /// :raises NoEdgeBetweenNodes: When there is no edge between nodes
+    #[text_signature = "(self, source, target, edge /)"]
+    pub fn update_edge_by_index(
+        &mut self,
+        edge_index: usize,
+        edge: PyObject,
+    ) -> PyResult<()> {
+        match self.graph.edge_weight_mut(EdgeIndex::new(edge_index)) {
+            Some(data) => *data = edge,
+            None => {
+                return Err(PyIndexError::new_err("No edge found for index"))
+            }
+        };
+        Ok(())
+    }
+
     /// Return the node data for a given node index
     ///
     /// :param int node: The index for the node

src/graph.rs

@@ -416,6 +416,60 @@ impl PyGraph {
         Ok(data)
     }
 
+    /// Update an edge's weight/payload in place
+    ///
+    /// If there are parallel edges in the graph only one edge will be updated.
+    /// if you need to update a specific edge or need to ensure all parallel
+    /// edges get updated you should use
+    /// :meth:`~retworkx.PyGraph.update_edge_by_index` instead.
+    ///
+    /// :param int source: The index for the first node
+    /// :param int target: The index for the second node
+    ///
+    /// :raises NoEdgeBetweenNodes: When there is no edge between nodes
+    #[text_signature = "(self, source, target, edge /)"]
+    pub fn update_edge(
+        &mut self,
+        source: usize,
+        target: usize,
+        edge: PyObject,
+    ) -> PyResult<()> {
+        let index_a = NodeIndex::new(source);
+        let index_b = NodeIndex::new(target);
+        let edge_index = match self.graph.find_edge(index_a, index_b) {
+            Some(edge_index) => edge_index,
+            None => {
+                return Err(NoEdgeBetweenNodes::new_err(
+                    "No edge found between nodes",
+                ))
+            }
+        };
+        let data = self.graph.edge_weight_mut(edge_index).unwrap();
+        *data = edge;
+        Ok(())
+    }
+
+    /// Update an edge's weight/data payload in place by the edge index
+    ///
+    /// :param int edge_index: The index for the edge
+    /// :param object edge: The data payload/weight to update the edge with
+    ///
+    /// :raises NoEdgeBetweenNodes: When there is no edge between nodes
+    #[text_signature = "(self, source, target, edge /)"]
+    pub fn update_edge_by_index(
+        &mut self,
+        edge_index: usize,
+        edge: PyObject,
+    ) -> PyResult<()> {
+        match self.graph.edge_weight_mut(EdgeIndex::new(edge_index)) {
+            Some(data) => *data = edge,
+            None => {
+                return Err(PyIndexError::new_err("No edge found for index"))
+            }
+        };
+        Ok(())
+    }
+
     /// Return the node data for a given node index
     ///
     /// :param int node: The index for the node

tests/graph/test_edges.py

@@ -42,6 +42,45 @@ def test_no_edge(self):
         self.assertRaises(retworkx.NoEdgeBetweenNodes, graph.get_edge_data,
                           node_a, node_b)
 
+    def test_update_edge(self):
+        graph = retworkx.PyGraph()
+        node_a = graph.add_node('a')
+        node_b = graph.add_node('b')
+        graph.add_edge(node_a, node_b, 'not edgy')
+        graph.update_edge(node_a, node_b, 'Edgy')
+        self.assertEqual([(0, 1, 'Edgy')], graph.weighted_edge_list())
+
+    def test_update_edge_no_edge(self):
+        graph = retworkx.PyGraph()
+        node_a = graph.add_node('a')
+        node_b = graph.add_node('b')
+        self.assertRaises(retworkx.NoEdgeBetweenNodes, graph.update_edge,
+                          node_a, node_b, None)
+
+    def test_update_edge_by_index(self):
+        graph = retworkx.PyGraph()
+        node_a = graph.add_node('a')
+        node_b = graph.add_node('b')
+        edge_index = graph.add_edge(node_a, node_b, 'not edgy')
+        graph.update_edge_by_index(edge_index, 'Edgy')
+        self.assertEqual([(0, 1, 'Edgy')], graph.weighted_edge_list())
+
+    def test_update_edge_invalid_index(self):
+        graph = retworkx.PyGraph()
+        graph.add_node('a')
+        graph.add_node('b')
+        self.assertRaises(IndexError, graph.update_edge_by_index, 0, None)
+
+    def test_update_edge_parallel_edges(self):
+        graph = retworkx.PyGraph()
+        node_a = graph.add_node('a')
+        node_b = graph.add_node('b')
+        graph.add_edge(node_a, node_b, 'not edgy')
+        edge_index = graph.add_edge(node_a, node_b, 'not edgy')
+        graph.update_edge_by_index(edge_index, 'Edgy')
+        self.assertEqual([(0, 1, 'not edgy'), (0, 1, 'Edgy')],
+                         list(graph.weighted_edge_list()))
+
     def test_no_edge_get_all_edge_data(self):
         graph = retworkx.PyGraph()
         node_a = graph.add_node('a')

tests/test_edges.py

@@ -40,6 +40,43 @@ def test_no_edge(self):
         self.assertRaises(retworkx.NoEdgeBetweenNodes, dag.get_edge_data,
                           node_a, node_b)
 
+    def test_update_edge(self):
+        dag = retworkx.PyDAG()
+        node_a = dag.add_node('a')
+        node_b = dag.add_child(node_a, 'b', 'not edgy')
+        dag.update_edge(node_a, node_b, 'Edgy')
+        self.assertEqual([(0, 1, 'Edgy')], dag.weighted_edge_list())
+
+    def test_update_edge_no_edge(self):
+        dag = retworkx.PyDAG()
+        node_a = dag.add_node('a')
+        node_b = dag.add_node('b')
+        self.assertRaises(retworkx.NoEdgeBetweenNodes, dag.update_edge,
+                          node_a, node_b, None)
+
+    def test_update_edge_by_index(self):
+        dag = retworkx.PyDAG()
+        node_a = dag.add_node('a')
+        dag.add_child(node_a, 'b', 'not edgy')
+        dag.update_edge_by_index(0, 'Edgy')
+        self.assertEqual([(0, 1, 'Edgy')], dag.weighted_edge_list())
+
+    def test_update_edge_invalid_index(self):
+        dag = retworkx.PyDAG()
+        dag.add_node('a')
+        dag.add_node('b')
+        self.assertRaises(IndexError, dag.update_edge_by_index, 0, None)
+
+    def test_update_edge_parallel_edges(self):
+        graph = retworkx.PyDiGraph()
+        node_a = graph.add_node('a')
+        node_b = graph.add_node('b')
+        graph.add_edge(node_a, node_b, 'not edgy')
+        edge_index = graph.add_edge(node_a, node_b, 'not edgy')
+        graph.update_edge_by_index(edge_index, 'Edgy')
+        self.assertEqual([(0, 1, 'not edgy'), (0, 1, 'Edgy')],
+                         list(graph.weighted_edge_list()))
+
     def test_has_edge(self):
         dag = retworkx.PyDAG()
         node_a = dag.add_node('a')