Move lollipop_graph generator to rustworkx-core

rustworkx-core/src/generators/lollipop_graph.rs

@@ -0,0 +1,200 @@
+// 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 std::hash::Hash;
+
+use petgraph::data::{Build, Create};
+use petgraph::visit::{Data, GraphProp, NodeIndexable};
+
+use super::utils::pairwise;
+use super::InvalidInputError;
+
+/// Generate a lollipop graph
+///
+/// Arguments:
+///
+/// Generate an undirected lollipop graph where a mesh graph is connected to a
+/// path.
+///
+/// If neither `num_path_nodes` nor `path_weights` (both described
+/// below) are specified then this is equivalent to
+/// :func:`~rustworkx.generators.mesh_graph`
+///
+/// * `num_mesh_nodes` - The number of nodes to generate the mesh graph
+///     with. Node weights will be None if this is specified. If both
+///     `num_mesh_nodes` and `mesh_weights` are set this will be ignored and
+///     `mesh_weights` will be used.
+/// * `num_path_nodes` - The number of nodes to generate the path
+///     with. Node weights will be None if this is specified. If both
+///     `num_path_nodes` and `path_weights` are set this will be ignored and
+///     `path_weights` will be used.
+/// * `mesh_weights` - A list of node weights for the mesh graph. If both
+///     `num_mesh_nodes` and `mesh_weights` are set `num_mesh_nodes` will
+///     be ignored and `mesh_weights` will be used.
+/// * `path_weights` - A list of node weights for the path. If both
+///     `num_path_nodes` and `path_weights` are set `num_path_nodes` will
+///     be ignored and `path_weights` will be used.
+/// * `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.
+///
+/// # Example
+/// ```rust
+/// use rustworkx_core::petgraph;
+/// use rustworkx_core::generators::lollipop_graph;
+/// use rustworkx_core::petgraph::visit::EdgeRef;
+///
+/// let expected_edge_list = vec![
+///     (0, 1),
+///     (0, 2),
+///     (0, 3),
+///     (1, 2),
+///     (1, 3),
+///     (2, 3),
+///     (3, 4),
+///     (4, 5),
+///     (5, 6),
+/// ];
+/// let g: petgraph::graph::UnGraph<(), ()> = lollipop_graph(
+///     Some(4),
+///     Some(3),
+///     None,
+///     None,
+///     || {()},
+///     || {()},
+/// ).unwrap();
+/// assert_eq!(
+///     expected_edge_list,
+///     g.edge_references()
+///         .map(|edge| (edge.source().index(), edge.target().index()))
+///         .collect::<Vec<(usize, usize)>>(),
+/// )
+/// ```
+pub fn lollipop_graph<G, T, F, H, M>(
+    num_mesh_nodes: Option<usize>,
+    num_path_nodes: Option<usize>,
+    mesh_weights: Option<Vec<T>>,
+    path_weights: Option<Vec<T>>,
+    mut default_node_weight: F,
+    mut default_edge_weight: H,
+) -> Result<G, InvalidInputError>
+where
+    G: Build + Create + Data<NodeWeight = T, EdgeWeight = M> + NodeIndexable + GraphProp,
+    F: FnMut() -> T,
+    H: FnMut() -> M,
+    G::NodeId: Eq + Hash,
+{
+    if num_mesh_nodes.is_none() && mesh_weights.is_none() {
+        return Err(InvalidInputError {});
+    }
+    let num_nodes: usize;
+    if let Some(mesh_nodes) = num_mesh_nodes {
+        num_nodes = mesh_nodes;
+    } else {
+        num_nodes = mesh_weights.as_ref().unwrap().len();
+    }
+    let num_edges = (num_nodes * (num_nodes - 1)) / 2;
+    let mut graph = G::with_capacity(num_nodes, num_edges);
+
+    let mesh_nodes: Vec<G::NodeId> = match mesh_weights {
+        Some(mesh_weights) => mesh_weights
+            .into_iter()
+            .map(|weight| graph.add_node(weight))
+            .collect(),
+        None => (0..num_mesh_nodes.unwrap())
+            .map(|_| graph.add_node(default_node_weight()))
+            .collect(),
+    };
+
+    let meshlen = mesh_nodes.len();
+    for i in 0..meshlen - 1 {
+        for j in i + 1..meshlen {
+            graph.add_edge(mesh_nodes[i], mesh_nodes[j], default_edge_weight());
+        }
+    }
+    let path_nodes: Vec<G::NodeId> = match path_weights {
+        Some(path_weights) => path_weights
+            .into_iter()
+            .map(|weight| graph.add_node(weight))
+            .collect(),
+        None => {
+            if let Some(num_path) = num_path_nodes {
+                (0..num_path)
+                    .map(|_| graph.add_node(default_node_weight()))
+                    .collect()
+            } else {
+                vec![]
+            }
+        }
+    };
+    let pathlen = path_nodes.len();
+    if pathlen > 0 {
+        graph.add_edge(
+            graph.from_index(meshlen - 1),
+            graph.from_index(meshlen),
+            default_edge_weight(),
+        );
+        for (node_a, node_b) in pairwise(path_nodes) {
+            match node_a {
+                Some(node_a) => graph.add_edge(node_a, node_b, default_edge_weight()),
+                None => continue,
+            };
+        }
+    }
+    Ok(graph)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::generators::lollipop_graph;
+    use crate::generators::InvalidInputError;
+    use crate::petgraph::visit::EdgeRef;
+    #[test]
+    fn test_lollipop_mesh_path() {
+        let expected_edge_list = vec![
+            (0, 1),
+            (0, 2),
+            (0, 3),
+            (1, 2),
+            (1, 3),
+            (2, 3),
+            (3, 4),
+            (4, 5),
+            (5, 6),
+        ];
+        let g: petgraph::graph::UnGraph<(), ()> =
+            lollipop_graph(Some(4), Some(3), None, None, || (), || ()).unwrap();
+        assert_eq!(
+            expected_edge_list,
+            g.edge_references()
+                .map(|edge| (edge.source().index(), edge.target().index()))
+                .collect::<Vec<(usize, usize)>>(),
+        );
+    }
+
+    #[test]
+    fn test_lollipop_none_mesh() {
+        match lollipop_graph::<petgraph::graph::UnGraph<(), ()>, (), _, _, ()>(
+            None,
+            None,
+            None,
+            None,
+            || (),
+            || (),
+        ) {
+            Ok(_) => panic!("Returned a non-error"),
+            Err(e) => assert_eq!(e, InvalidInputError),
+        };
+    }
+}

rustworkx-core/src/generators/mod.rs

@@ -19,6 +19,7 @@ mod grid_graph;
 mod heavy_hex_graph;
 mod heavy_square_graph;
 mod hexagonal_lattice_graph;
+mod lollipop_graph;
 mod path_graph;
 mod petersen_graph;
 mod star_graph;
@@ -47,6 +48,7 @@ pub use grid_graph::grid_graph;
 pub use heavy_hex_graph::heavy_hex_graph;
 pub use heavy_square_graph::heavy_square_graph;
 pub use hexagonal_lattice_graph::hexagonal_lattice_graph;
+pub use lollipop_graph::lollipop_graph;
 pub use path_graph::path_graph;
 pub use petersen_graph::petersen_graph;
 pub use star_graph::star_graph;

src/generators.rs

@@ -1332,37 +1332,28 @@ pub fn lollipop_graph(
     path_weights: Option<Vec<PyObject>>,
     multigraph: bool,
 ) -> PyResult<graph::PyGraph> {
-    let mut graph = mesh_graph(py, num_mesh_nodes, mesh_weights, multigraph)?;
-    if num_path_nodes.is_none() && path_weights.is_none() {
-        return Ok(graph);
-    }
-    let meshlen = graph.num_nodes();
-
-    let path_nodes: Vec<NodeIndex> = match path_weights {
-        Some(path_weights) => path_weights
-            .into_iter()
-            .map(|node| graph.graph.add_node(node))
-            .collect(),
-        None => (0..num_path_nodes.unwrap())
-            .map(|_| graph.graph.add_node(py.None()))
-            .collect(),
-    };
-
-    let pathlen = path_nodes.len();
-    if pathlen > 0 {
-        graph.graph.add_edge(
-            NodeIndex::new(meshlen - 1),
-            NodeIndex::new(meshlen),
-            py.None(),
-        );
-        for (node_a, node_b) in pairwise(path_nodes) {
-            match node_a {
-                Some(node_a) => graph.graph.add_edge(node_a, node_b, py.None()),
-                None => continue,
-            };
+    let default_fn = || py.None();
+    let graph: StablePyGraph<Undirected> = match core_generators::lollipop_graph(
+        num_mesh_nodes,
+        num_path_nodes,
+        mesh_weights,
+        path_weights,
+        default_fn,
+        default_fn,
+    ) {
+        Ok(graph) => graph,
+        Err(_) => {
+            return Err(PyIndexError::new_err(
+                "num_nodes and weights list not specified",
+            ))
         }
-    }
-    Ok(graph)
+    };
+    Ok(graph::PyGraph {
+        graph,
+        node_removed: false,
+        multigraph,
+        attrs: py.None(),
+    })
 }
 
 /// Generate a generalized Petersen graph :math:`G(n, k)` with :math:`2n`