Adding generator for creating graphs for testing.

Author: root-11

graph/generate.py

@@ -0,0 +1,38 @@
+import itertools
+from graph.core import Graph
+
+
+def binary_tree(levels):
+    """
+    Generates a binary tree with the given number of levels.
+    """
+    if not isinstance(levels, int):
+        raise TypeError(f"Expected int, not {type(levels)}")
+    g = Graph()
+    for i in range(levels):
+        g.add_edge(i, 2 * i + 1)
+        g.add_edge(i, 2 * i + 2)
+    return g
+
+
+def grid(length, width, bidirectional=False):
+    """
+    Generates a grid with the given length and width.
+    """
+    if not isinstance(length, int):
+        raise TypeError(f"Expected int, not {type(length)}")
+    if not isinstance(width, int):
+        raise TypeError(f"Expected int, not {type(width)}")
+    g = Graph()
+    node_index = {}
+    c = itertools.count(start=1)
+    for i in range(length):  # i is the row
+        for j in range(width):  # j is the column
+            node_index[(i, j)] = next(c)
+            if i > 0:
+                a, b = node_index[(i, j)], node_index[(i - 1, j)]
+                g.add_edge(b, a, bidirectional=bidirectional)
+            if j > 0:
+                a, b = node_index[(i, j)], node_index[(i, j - 1)]
+                g.add_edge(b, a, bidirectional=bidirectional)
+    return g

tests/test_generate.py

@@ -0,0 +1,26 @@
+from graph.core import Graph
+from graph.generate import binary_tree, grid
+
+from tests.test_graph import graph3x3, graph4x4, graph5x5
+
+
+def test_generate_grid_3x3():
+    g = grid(3, 3)
+    expected = graph3x3()
+    assert g == expected
+
+
+def test_generate_grid_4x4_bidirectional():
+    g = grid(4, 4, bidirectional=True)
+    expected = graph4x4()
+    assert g == expected
+
+def test_generate_grid_5x5_bidirectional():
+    g = grid(5, 5, bidirectional=True)
+    expected = graph5x5()
+    assert g == expected
+
+def test_generate_binary_tree_3_levels():
+    g = binary_tree(3)
+    expected = Graph(from_list=[[0, 1], [0, 2], [1, 3], [1, 4], [2, 5], [2, 6]])
+    assert g == expected