From 15be3442423ebc7b4dbb154f737c48bb7affbdae Mon Sep 17 00:00:00 2001
From: Seungwook-Woo <78723913+Seungwook-Woo@users.noreply.github.com>
Date: Wed, 19 Oct 2022 20:48:28 +0900
Subject: [PATCH] Added ladder graph generators and test_ladder

---
 src/generators.rs                             | 88 +++++++++++++++++++
 .../rustworkx_tests/generators/test_ladder.py | 26 ++++++
 2 files changed, 114 insertions(+)
 create mode 100644 tests/rustworkx_tests/generators/test_ladder.py

diff --git a/src/generators.rs b/src/generators.rs
index e497c674a4..e7701474f7 100644
--- a/src/generators.rs
+++ b/src/generators.rs
@@ -2420,6 +2420,93 @@ pub fn barbell_graph(
     })
 }
 
+/// Generate a ladder graph
+///
+/// :param int num_node: The number of nodes to generate the graph with. Node
+///     weights will be None if this is specified. If both ``num_node`` and
+///     ``weights`` are set this will be ignored and ``weights``will be used.
+/// :param list weights: A list of node weights. 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 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 ladder graph
+/// :rtype: PyGraph
+/// :raises IndexError: If neither ``num_nodes`` or ``weights`` are specified
+///
+/// .. jupyter-execute::
+///
+///   import rustworkx.generators
+///   from rustworkx.visualization import mpl_draw
+///
+///   graph = rustworkx.generators.ladder_graph(5)
+///   mpl_draw(graph)
+///
+#[pyfunction(multigraph = true)]
+#[pyo3(text_signature = "(/, num_nodes=None, weights=None, multigraph=True")]
+pub fn ladder_graph(
+    py: Python,
+    num_nodes: Option<usize>,
+    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",
+        ));
+    } else if num_nodes.is_none() && weights.as_ref().unwrap().len() % 2 == 1 {
+        return Err(PyIndexError::new_err(
+            "length of weights must be even numbers",
+        ));
+    }
+    let node_len = if weights.is_none() {
+        num_nodes.unwrap()
+    } else {
+        weights.as_ref().unwrap().len() / 2
+    };
+    
+    if node_len == 0 {
+        return Ok(graph::PyGraph {
+            graph: StablePyGraph::<Undirected>::default(),
+            node_removed: false,
+            multigraph,
+            attrs: py.None(),
+        });
+    }
+    let num_edges = 3 * node_len - 2;
+    let mut graph = StablePyGraph::<Undirected>::with_capacity(node_len * 2, num_edges);
+    match weights {
+        Some(weights) => {
+            for weight in weights {
+                graph.add_node(weight);
+            }
+        }
+        None => {
+            (0..(node_len * 2)).for_each(|_| {
+                graph.add_node(py.None());
+            });
+        }
+    };
+
+    for rail_a in 0..node_len - 1 {
+        graph.add_edge(NodeIndex::new(rail_a), NodeIndex::new(rail_a+1), py.None());
+        let rail_b = rail_a + node_len;
+        graph.add_edge(NodeIndex::new(rail_b), NodeIndex::new(rail_b+1), py.None());
+    }
+    for rung_a in 0..node_len {
+        graph.add_edge(NodeIndex::new(rung_a), NodeIndex::new(rung_a+node_len), py.None());
+    }
+    
+    Ok(graph::PyGraph {
+        graph,
+        node_removed: false,
+        multigraph,
+        attrs: py.None(),
+    })
+}
+
 #[pymodule]
 pub fn generators(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(cycle_graph))?;
@@ -2444,5 +2531,6 @@ pub fn generators(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(full_rary_tree))?;
     m.add_wrapped(wrap_pyfunction!(generalized_petersen_graph))?;
     m.add_wrapped(wrap_pyfunction!(barbell_graph))?;
+    m.add_wrapped(wrap_pyfunction!(ladder_graph))?;
     Ok(())
 }
diff --git a/tests/rustworkx_tests/generators/test_ladder.py b/tests/rustworkx_tests/generators/test_ladder.py
new file mode 100644
index 0000000000..7c05b93104
--- /dev/null
+++ b/tests/rustworkx_tests/generators/test_ladder.py
@@ -0,0 +1,26 @@
+# Licensed under the Apache License, 
+
+import unittest
+
+import rustworkx
+
+
+class TestLadderGraph(unittest.TestCase):
+    def test_ladder_graph(self):
+        graph = rustworkx.generators.ladder_graph(20)
+        self.assertEqual(len(graph), 40)
+        self.assertEqual(len(graph.edges()), 58)
+    
+    def test_ladder_graph_weights(self):
+        graph = rustworkx.generators.ladder_graph(weights=list(range(40)))
+        self.assertEqual(len(graph), 40)
+        self.assertEqual([x for x in range(40)], graph.nodes())
+        self.assertEqual(len(graph.edges()), 58)
+
+    def test_ladder_no_weights_or_num(self):
+        with self.assertRaises(IndexError):
+            rustworkx.generators.ladder_graph()
+
+    def test_zero_length_ladder_graph(self):
+        graph = rustworkx.generators.ladder_graph(0)
+        self.assertEqual(0, len(graph))
\ No newline at end of file