Add substitute_node_with_subgraph method to PyDiGraph

releasenotes/notes/add-subgraph-replace-ec54792df1d641fb.yaml

@@ -0,0 +1,34 @@
+---
+features:
+  - |
+    A new method, :meth:`~retworkx.PyDiGraph.substitute_node_with_subgraph`,
+    to the :class:`~retworkx.PyDiGraph` class. This method is used to replace a
+    node in a :class:`~retworkx.PyDiGraph` object with another
+    :class:`~retwork.PyDiGraph` object. For example, first creating a new
+    graph:
+
+    .. jupyter-execute::
+
+      import retworkx
+      from retworkx.visualization import mpl_draw
+
+      original_graph = retworkx.generators.directed_path_graph(5)
+      mpl_draw(original_graph)
+
+    then create another graph to use in place of a node:
+
+    .. jupyter-execute::
+
+      other_graph = retworkx.generators.directed_star_graph(25)
+      mpl_draw(other_graph)
+
+    finally replace a node in the original graph with the second graph:
+
+    .. jupyter-execute::
+
+      def edge_map_fn(_source, _target, _weight):
+          return 0
+
+      node_mapping = original_graph.substitute_node_with_subgraph(2, other_graph, edge_map_fn)
+      print(node_mapping)
+      mpl_draw(original_graph)

src/digraph.rs

@@ -47,7 +47,7 @@ use petgraph::visit::{
 
 use super::dot_utils::build_dot;
 use super::iterators::{
-    EdgeIndexMap, EdgeIndices, EdgeList, NodeIndices, WeightedEdgeList,
+    EdgeIndexMap, EdgeIndices, EdgeList, NodeIndices, NodeMap, WeightedEdgeList,
 };
 use super::{
     is_directed_acyclic_graph, DAGHasCycle, DAGWouldCycle, NoEdgeBetweenNodes,
@@ -2203,6 +2203,153 @@ impl PyDiGraph {
         Ok(out_dict.into())
     }
 
+    /// Substitute a node with a PyDigraph object
+    ///
+    /// :param int node: The node to replace with the PyDiGraph object
+    /// :param PyDiGraph other: The other graph to replace ``node`` with
+    /// :param callable edge_map_fn: A callable object that will take 3 position
+    ///     parameters, ``(source, target, weight)`` to represent an edge either to
+    ///     or from ``node`` in this graph. The expected return value from this
+    ///     callable is the node index of the node in ``other`` that an edge should
+    ///     be to/from. If None is returned, that edge will be skipped and not
+    ///     be copied.
+    /// :param callable node_filter: An optional callable object that when used
+    ///     will receive a node's payload object from ``other`` and return
+    ///     ``True`` if that node is to be included in the graph or not.
+    /// :param callable edge_weight_map: An optional callable object that when
+    ///     used will receive an edge's weight/data payload from ``other`` and
+    ///     will return an object to use as the weight for a newly created edge
+    ///     after the edge is mapped from ``other``. If not specified the weight
+    ///     from the edge in ``other`` will be copied by reference and used.
+    ///
+    /// :returns: A mapping of node indices in ``other`` to the equivalent node
+    ///     in this graph.
+    /// :rtype: NodeMap
+    #[text_signature = "(self, node, other, edge_map_fn, /, node_filter=None, edge_weight_map=None)"]
+    fn substitute_node_with_subgraph(
+        &mut self,
+        py: Python,
+        node: usize,
+        other: &PyDiGraph,
+        edge_map_fn: PyObject,
+        node_filter: Option<PyObject>,
+        edge_weight_map: Option<PyObject>,
+    ) -> PyResult<NodeMap> {
+        let weight_map_fn = |obj: &PyObject,
+                             weight_fn: &Option<PyObject>|
+         -> PyResult<PyObject> {
+            match weight_fn {
+                Some(weight_fn) => weight_fn.call1(py, (obj,)),
+                None => Ok(obj.clone_ref(py)),
+            }
+        };
+        let map_fn = |source: usize,
+                      target: usize,
+                      weight: &PyObject|
+         -> PyResult<Option<usize>> {
+            let res = edge_map_fn.call1(py, (source, target, weight))?;
+            res.extract(py)
+        };
+        let filter_fn =
+            |obj: &PyObject, filter_fn: &Option<PyObject>| -> PyResult<bool> {
+                match filter_fn {
+                    Some(filter) => {
+                        let res = filter.call1(py, (obj,))?;
+                        res.extract(py)
+                    }
+                    None => Ok(true),
+                }
+            };
+        let node_index: NodeIndex = NodeIndex::new(node);
+        if self.graph.node_weight(node_index).is_none() {
+            return Err(PyIndexError::new_err(format!(
+                "Specified node {} is not in this graph",
+                node
+            )));
+        }
+        // Copy nodes from other to self
+        let mut out_map: HashMap<usize, usize> =
+            HashMap::with_capacity(other.node_count());
+        for node in other.graph.node_indices() {
+            let node_weight = other[node].clone_ref(py);
+            if !filter_fn(&node_weight, &node_filter)? {
+                continue;
+            }
+            let new_index = self.graph.add_node(node_weight);
+            out_map.insert(node.index(), new_index.index());
+        }
+        // If no nodes are copied bail here since there is nothing left
+        // to do.
+        if out_map.is_empty() {
+            self.graph.remove_node(node_index);
+            // Return a new empty map to clear allocation from out_map
+            return Ok(NodeMap {
+                node_map: HashMap::new(),
+            });
+        }
+        // Copy edges from other to self
+        for edge in other.graph.edge_references().filter(|edge| {
+            out_map.contains_key(&edge.target().index())
+                && out_map.contains_key(&edge.source().index())
+        }) {
+            self._add_edge(
+                NodeIndex::new(out_map[&edge.source().index()]),
+                NodeIndex::new(out_map[&edge.target().index()]),
+                weight_map_fn(edge.weight(), &edge_weight_map)?,
+            )?;
+        }
+        // Add edges to/from node to nodes in other
+        let in_edges: Vec<(NodeIndex, NodeIndex, PyObject)> = self
+            .graph
+            .edges_directed(node_index, petgraph::Direction::Incoming)
+            .map(|edge| {
+                (edge.source(), edge.target(), edge.weight().clone_ref(py))
+            })
+            .collect();
+        let out_edges: Vec<(NodeIndex, NodeIndex, PyObject)> = self
+            .graph
+            .edges_directed(node_index, petgraph::Direction::Outgoing)
+            .map(|edge| {
+                (edge.source(), edge.target(), edge.weight().clone_ref(py))
+            })
+            .collect();
+        for (source, target, weight) in in_edges {
+            let old_index = map_fn(source.index(), target.index(), &weight)?;
+            let target_out = match old_index {
+                Some(old_index) => match out_map.get(&old_index) {
+                    Some(new_index) => NodeIndex::new(*new_index),
+                    None => {
+                        return Err(PyIndexError::new_err(format!(
+                            "No mapped index {} found",
+                            old_index
+                        )))
+                    }
+                },
+                None => continue,
+            };
+            self._add_edge(source, target_out, weight)?;
+        }
+        for (source, target, weight) in out_edges {
+            let old_index = map_fn(source.index(), target.index(), &weight)?;
+            let source_out = match old_index {
+                Some(old_index) => match out_map.get(&old_index) {
+                    Some(new_index) => NodeIndex::new(*new_index),
+                    None => {
+                        return Err(PyIndexError::new_err(format!(
+                            "No mapped index {} found",
+                            old_index
+                        )))
+                    }
+                },
+                None => continue,
+            };
+            self._add_edge(source_out, target, weight)?;
+        }
+        // Remove node
+        self.graph.remove_node(node_index);
+        Ok(NodeMap { node_map: out_map })
+    }
+
     /// Return a new PyDiGraph object for a subgraph of this graph
     ///
     /// :param list nodes: A list of node indices to generate the subgraph