Move star_graph generator to rustworkx-core

rustworkx-core/src/generators/mod.rs

@@ -13,6 +13,7 @@
 //! This module contains generator functions for building graphs
 
 mod cycle_graph;
+mod star_graph;
 
 mod utils;
 
@@ -32,3 +33,4 @@ impl fmt::Display for InvalidInputError {
 }
 
 pub use cycle_graph::cycle_graph;
+pub use star_graph::star_graph;

rustworkx-core/src/generators/star_graph.rs

@@ -0,0 +1,179 @@
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations
+// under the License.
+
+use petgraph::data::{Build, Create};
+use petgraph::visit::{Data, NodeIndexable};
+
+use super::utils::get_num_nodes;
+use super::InvalidInputError;
+
+/// Generate a star graph
+///
+/// Arguments:
+///
+/// * `num_nodes` - The number of nodes to create a path graph for. Either this or
+///     `weights` must be specified. If both this and `weights` are specified, weights
+///     will take priorty and this argument will be ignored
+/// * `weights` - A `Vec` of node weight objects.
+/// * `default_node_weight` - A callable that will return the weight to use
+///     for newly created nodes. This is ignored if `weights` is specified,
+///     as the weights from that argument will be used instead.
+/// * `default_edge_weight` - A callable that will return the weight object
+///     to use for newly created edges.
+/// *  inward: If set `true` the nodes will be directed towards the
+///     center node. This parameter is ignored if `bidirectional` is set to
+///     `true`.
+/// * `bidirectional` - Whether edges are added bidirectionally, if set to
+///     `true` then for any edge `(u, v)` an edge `(v, u)` will also be added.
+///     If the graph is undirected this will result in a parallel edge.
+
+/// # Example
+/// ```rust
+/// use rustworkx_core::petgraph;
+/// use rustworkx_core::generators::star_graph;
+/// use rustworkx_core::petgraph::visit::EdgeRef;
+///
+/// let g: petgraph::graph::UnGraph<(), ()> = star_graph(
+///     Some(4),
+///     None,
+///     || {()},
+///     || {()},
+///     false,
+///     false
+/// ).unwrap();
+/// assert_eq!(
+///     vec![(0, 1), (0, 2), (0, 3)],
+///     g.edge_references()
+///         .map(|edge| (edge.source().index(), edge.target().index()))
+///         .collect::<Vec<(usize, usize)>>(),
+/// )
+/// ```
+pub fn star_graph<G, T, F, H, M>(
+    num_nodes: Option<usize>,
+    weights: Option<Vec<T>>,
+    mut default_node_weight: F,
+    mut default_edge_weight: H,
+    inward: bool,
+    bidirectional: bool,
+) -> Result<G, InvalidInputError>
+where
+    G: Build + Create + Data<NodeWeight = T, EdgeWeight = M> + NodeIndexable,
+    F: FnMut() -> T,
+    H: FnMut() -> M,
+{
+    if weights.is_none() && num_nodes.is_none() {
+        return Err(InvalidInputError {});
+    }
+    let node_len = get_num_nodes(&num_nodes, &weights);
+    let num_edges = if bidirectional {
+        2 * node_len
+    } else {
+        node_len
+    };
+    let mut graph = G::with_capacity(node_len, num_edges);
+    if node_len == 0 {
+        return Ok(graph);
+    }
+
+    match weights {
+        Some(weights) => {
+            for weight in weights {
+                graph.add_node(weight);
+            }
+        }
+        None => {
+            for _ in 0..node_len {
+                graph.add_node(default_node_weight());
+            }
+        }
+    };
+    let zero_index = graph.from_index(0);
+    for a in 1..node_len {
+        let node = graph.from_index(a);
+        if bidirectional {
+            graph.add_edge(node, zero_index, default_edge_weight());
+            graph.add_edge(zero_index, node, default_edge_weight());
+        } else if inward {
+            graph.add_edge(node, zero_index, default_edge_weight());
+        } else {
+            graph.add_edge(zero_index, node, default_edge_weight());
+        }
+    }
+    Ok(graph)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::generators::star_graph;
+    use crate::generators::InvalidInputError;
+    use crate::petgraph;
+    use crate::petgraph::visit::EdgeRef;
+
+    #[test]
+    fn test_with_weights() {
+        let g: petgraph::graph::UnGraph<usize, ()> =
+            star_graph(None, Some(vec![0, 1, 2, 3]), || 4, || (), false, false).unwrap();
+        assert_eq!(
+            vec![(0, 1), (0, 2), (0, 3)],
+            g.edge_references()
+                .map(|edge| (edge.source().index(), edge.target().index()))
+                .collect::<Vec<(usize, usize)>>(),
+        );
+        assert_eq!(
+            vec![0, 1, 2, 3],
+            g.node_weights().copied().collect::<Vec<usize>>(),
+        );
+    }
+
+    #[test]
+    fn test_with_weights_inward() {
+        let g: petgraph::graph::UnGraph<usize, ()> =
+            star_graph(None, Some(vec![0, 1, 2, 3]), || 4, || (), true, false).unwrap();
+        assert_eq!(
+            vec![(1, 0), (2, 0), (3, 0)],
+            g.edge_references()
+                .map(|edge| (edge.source().index(), edge.target().index()))
+                .collect::<Vec<(usize, usize)>>(),
+        );
+        assert_eq!(
+            vec![0, 1, 2, 3],
+            g.node_weights().copied().collect::<Vec<usize>>(),
+        );
+    }
+
+    #[test]
+    fn test_bidirectional() {
+        let g: petgraph::graph::DiGraph<(), ()> =
+            star_graph(Some(4), None, || (), || (), false, true).unwrap();
+        assert_eq!(
+            vec![(1, 0), (0, 1), (2, 0), (0, 2), (3, 0), (0, 3)],
+            g.edge_references()
+                .map(|edge| (edge.source().index(), edge.target().index()))
+                .collect::<Vec<(usize, usize)>>(),
+        );
+    }
+
+    #[test]
+    fn test_error() {
+        match star_graph::<petgraph::graph::DiGraph<(), ()>, (), _, _, ()>(
+            None,
+            None,
+            || (),
+            || (),
+            false,
+            false,
+        ) {
+            Ok(_) => panic!("Returned a non-error"),
+            Err(e) => assert_eq!(e, InvalidInputError),
+        };
+    }
+}

src/generators.rs

@@ -321,11 +321,11 @@ pub fn path_graph(
 /// :param list weights: A list of node weights, the first element in the list
 ///     will be the center node of the star graph. If both ``num_node`` and
 ///     ``weights`` are set this will be ignored and ``weights`` will be used.
-/// :param bool bidirectional: Adds edges in both directions between two nodes
-///     if set to ``True``. Default value is ``False``.
 /// :param bool inward: If set ``True`` the nodes will be directed towards the
 ///     center node. This parameter is ignored if ``bidirectional`` is set to
 ///     ``True``.
+/// :param bool bidirectional: Adds edges in both directions between two nodes
+///     if set to ``True``. Default value is ``False``.
 /// :param bool multigraph: When set to False the output
 ///     :class:`~rustworkx.PyDiGraph` object will not be not be a multigraph and
 ///     won't allow parallel edges to be added. Instead
@@ -363,53 +363,22 @@ pub fn directed_star_graph(
     bidirectional: bool,
     multigraph: bool,
 ) -> PyResult<digraph::PyDiGraph> {
-    if weights.is_none() && num_nodes.is_none() {
-        return Err(PyIndexError::new_err(
-            "num_nodes and weights list not specified",
-        ));
-    }
-    let node_len = get_num_nodes(&num_nodes, &weights);
-    if node_len == 0 {
-        return Ok(digraph::PyDiGraph {
-            graph: StablePyGraph::<Directed>::default(),
-            node_removed: false,
-            check_cycle: false,
-            cycle_state: algo::DfsSpace::default(),
-            multigraph,
-            attrs: py.None(),
-        });
-    }
-    let num_edges = if bidirectional {
-        (2 * node_len) - 2
-    } else {
-        node_len - 1
-    };
-    let mut graph = StablePyGraph::<Directed>::with_capacity(node_len, num_edges);
-    match weights {
-        Some(weights) => {
-            for weight in weights {
-                graph.add_node(weight);
-            }
-        }
-        None => {
-            (0..node_len).for_each(|_| {
-                graph.add_node(py.None());
-            });
+    let default_fn = || py.None();
+    let graph: StablePyGraph<Directed> = match core_generators::star_graph(
+        num_nodes,
+        weights,
+        default_fn,
+        default_fn,
+        inward,
+        bidirectional,
+    ) {
+        Ok(graph) => graph,
+        Err(_) => {
+            return Err(PyIndexError::new_err(
+                "num_nodes and weights list not specified",
+            ))
         }
     };
-    let zero_index = NodeIndex::new(0);
-    for node_index in 1..node_len {
-        //Add edges in both directions if bidirection is True
-        let node = NodeIndex::new(node_index);
-        if bidirectional {
-            graph.add_edge(node, zero_index, py.None());
-            graph.add_edge(zero_index, node, py.None());
-        } else if inward {
-            graph.add_edge(node, zero_index, py.None());
-        } else {
-            graph.add_edge(zero_index, node, py.None());
-        }
-    }
     Ok(digraph::PyDiGraph {
         graph,
         node_removed: false,
@@ -429,10 +398,11 @@ pub fn directed_star_graph(
 ///     will be the center node of the star graph. If both ``num_node`` and
 ///     ``weights`` are set this will be ignored and ``weights`` will be used.
 /// :param bool multigraph: When set to False the output
-///     :class:`~rustworkx.PyGraph` object will not be not be a multigraph and
-///     won't  allow parallel edges to be added. Instead
+///     :class:`~rustworkx.PyDiGraph` object will not be not be a multigraph and
+///     won't allow parallel edges to be added. Instead
 ///     calls which would create a parallel edge will update the existing edge.
 ///
+///
 /// :returns: The generated star graph
 /// :rtype: PyGraph
 /// :raises IndexError: If neither ``num_nodes`` or ``weights`` are specified
@@ -453,37 +423,17 @@ pub fn star_graph(
     weights: Option<Vec<PyObject>>,
     multigraph: bool,
 ) -> PyResult<graph::PyGraph> {
-    if weights.is_none() && num_nodes.is_none() {
-        return Err(PyIndexError::new_err(
-            "num_nodes and weights list not specified",
-        ));
-    }
-    let node_len = get_num_nodes(&num_nodes, &weights);
-    if node_len == 0 {
-        return Ok(graph::PyGraph {
-            graph: StablePyGraph::<Undirected>::default(),
-            node_removed: false,
-            multigraph,
-            attrs: py.None(),
-        });
-    }
-    let mut graph = StablePyGraph::<Undirected>::with_capacity(node_len, node_len - 1);
-    match weights {
-        Some(weights) => {
-            for weight in weights {
-                graph.add_node(weight);
+    let default_fn = || py.None();
+    let graph: StablePyGraph<Undirected> =
+        match core_generators::star_graph(num_nodes, weights, default_fn, default_fn, false, false)
+        {
+            Ok(graph) => graph,
+            Err(_) => {
+                return Err(PyIndexError::new_err(
+                    "num_nodes and weights list not specified",
+                ))
             }
-        }
-        None => {
-            (0..node_len).for_each(|_| {
-                graph.add_node(py.None());
-            });
-        }
-    };
-    let zero_index = NodeIndex::new(0);
-    for node in 1..node_len {
-        graph.add_edge(zero_index, NodeIndex::new(node), py.None());
-    }
+        };
     Ok(graph::PyGraph {
         graph,
         node_removed: false,