test_graphlib.py 8.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252
  1. import graphlib
  2. import os
  3. import unittest
  4. from test.support.script_helper import assert_python_ok
  5. class TestTopologicalSort(unittest.TestCase):
  6. def _test_graph(self, graph, expected):
  7. def static_order_with_groups(ts):
  8. ts.prepare()
  9. while ts.is_active():
  10. nodes = ts.get_ready()
  11. for node in nodes:
  12. ts.done(node)
  13. yield tuple(sorted(nodes))
  14. ts = graphlib.TopologicalSorter(graph)
  15. self.assertEqual(list(static_order_with_groups(ts)), list(expected))
  16. ts = graphlib.TopologicalSorter(graph)
  17. # need to be a bit careful comparing the result of ts.static_order and
  18. # expected, because the order within a group is dependent on set
  19. # iteration order
  20. it = iter(ts.static_order())
  21. for group in expected:
  22. tsgroup = {next(it) for element in group}
  23. self.assertEqual(set(group), tsgroup)
  24. def _assert_cycle(self, graph, cycle):
  25. ts = graphlib.TopologicalSorter()
  26. for node, dependson in graph.items():
  27. ts.add(node, *dependson)
  28. try:
  29. ts.prepare()
  30. except graphlib.CycleError as e:
  31. _, seq = e.args
  32. self.assertIn(" ".join(map(str, cycle)), " ".join(map(str, seq * 2)))
  33. else:
  34. raise
  35. def test_simple_cases(self):
  36. self._test_graph(
  37. {2: {11}, 9: {11, 8}, 10: {11, 3}, 11: {7, 5}, 8: {7, 3}},
  38. [(3, 5, 7), (8, 11), (2, 9, 10)],
  39. )
  40. self._test_graph({1: {}}, [(1,)])
  41. self._test_graph(
  42. {x: {x + 1} for x in range(10)}, [(x,) for x in range(10, -1, -1)]
  43. )
  44. self._test_graph(
  45. {2: {3}, 3: {4}, 4: {5}, 5: {1}, 11: {12}, 12: {13}, 13: {14}, 14: {15}},
  46. [(1, 15), (5, 14), (4, 13), (3, 12), (2, 11)],
  47. )
  48. self._test_graph(
  49. {
  50. 0: [1, 2],
  51. 1: [3],
  52. 2: [5, 6],
  53. 3: [4],
  54. 4: [9],
  55. 5: [3],
  56. 6: [7],
  57. 7: [8],
  58. 8: [4],
  59. 9: [],
  60. },
  61. [(9,), (4,), (3, 8), (1, 5, 7), (6,), (2,), (0,)],
  62. )
  63. self._test_graph({0: [1, 2], 1: [], 2: [3], 3: []}, [(1, 3), (2,), (0,)])
  64. self._test_graph(
  65. {0: [1, 2], 1: [], 2: [3], 3: [], 4: [5], 5: [6], 6: []},
  66. [(1, 3, 6), (2, 5), (0, 4)],
  67. )
  68. def test_no_dependencies(self):
  69. self._test_graph({1: {2}, 3: {4}, 5: {6}}, [(2, 4, 6), (1, 3, 5)])
  70. self._test_graph({1: set(), 3: set(), 5: set()}, [(1, 3, 5)])
  71. def test_the_node_multiple_times(self):
  72. # Test same node multiple times in dependencies
  73. self._test_graph({1: {2}, 3: {4}, 0: [2, 4, 4, 4, 4, 4]}, [(2, 4), (0, 1, 3)])
  74. # Test adding the same dependency multiple times
  75. ts = graphlib.TopologicalSorter()
  76. ts.add(1, 2)
  77. ts.add(1, 2)
  78. ts.add(1, 2)
  79. self.assertEqual([*ts.static_order()], [2, 1])
  80. def test_graph_with_iterables(self):
  81. dependson = (2 * x + 1 for x in range(5))
  82. ts = graphlib.TopologicalSorter({0: dependson})
  83. self.assertEqual(list(ts.static_order()), [1, 3, 5, 7, 9, 0])
  84. def test_add_dependencies_for_same_node_incrementally(self):
  85. # Test same node multiple times
  86. ts = graphlib.TopologicalSorter()
  87. ts.add(1, 2)
  88. ts.add(1, 3)
  89. ts.add(1, 4)
  90. ts.add(1, 5)
  91. ts2 = graphlib.TopologicalSorter({1: {2, 3, 4, 5}})
  92. self.assertEqual([*ts.static_order()], [*ts2.static_order()])
  93. def test_empty(self):
  94. self._test_graph({}, [])
  95. def test_cycle(self):
  96. # Self cycle
  97. self._assert_cycle({1: {1}}, [1, 1])
  98. # Simple cycle
  99. self._assert_cycle({1: {2}, 2: {1}}, [1, 2, 1])
  100. # Indirect cycle
  101. self._assert_cycle({1: {2}, 2: {3}, 3: {1}}, [1, 3, 2, 1])
  102. # not all elements involved in a cycle
  103. self._assert_cycle({1: {2}, 2: {3}, 3: {1}, 5: {4}, 4: {6}}, [1, 3, 2, 1])
  104. # Multiple cycles
  105. self._assert_cycle({1: {2}, 2: {1}, 3: {4}, 4: {5}, 6: {7}, 7: {6}}, [1, 2, 1])
  106. # Cycle in the middle of the graph
  107. self._assert_cycle({1: {2}, 2: {3}, 3: {2, 4}, 4: {5}}, [3, 2])
  108. def test_calls_before_prepare(self):
  109. ts = graphlib.TopologicalSorter()
  110. with self.assertRaisesRegex(ValueError, r"prepare\(\) must be called first"):
  111. ts.get_ready()
  112. with self.assertRaisesRegex(ValueError, r"prepare\(\) must be called first"):
  113. ts.done(3)
  114. with self.assertRaisesRegex(ValueError, r"prepare\(\) must be called first"):
  115. ts.is_active()
  116. def test_prepare_multiple_times(self):
  117. ts = graphlib.TopologicalSorter()
  118. ts.prepare()
  119. with self.assertRaisesRegex(ValueError, r"cannot prepare\(\) more than once"):
  120. ts.prepare()
  121. def test_invalid_nodes_in_done(self):
  122. ts = graphlib.TopologicalSorter()
  123. ts.add(1, 2, 3, 4)
  124. ts.add(2, 3, 4)
  125. ts.prepare()
  126. ts.get_ready()
  127. with self.assertRaisesRegex(ValueError, "node 2 was not passed out"):
  128. ts.done(2)
  129. with self.assertRaisesRegex(ValueError, r"node 24 was not added using add\(\)"):
  130. ts.done(24)
  131. def test_done(self):
  132. ts = graphlib.TopologicalSorter()
  133. ts.add(1, 2, 3, 4)
  134. ts.add(2, 3)
  135. ts.prepare()
  136. self.assertEqual(ts.get_ready(), (3, 4))
  137. # If we don't mark anything as done, get_ready() returns nothing
  138. self.assertEqual(ts.get_ready(), ())
  139. ts.done(3)
  140. # Now 2 becomes available as 3 is done
  141. self.assertEqual(ts.get_ready(), (2,))
  142. self.assertEqual(ts.get_ready(), ())
  143. ts.done(4)
  144. ts.done(2)
  145. # Only 1 is missing
  146. self.assertEqual(ts.get_ready(), (1,))
  147. self.assertEqual(ts.get_ready(), ())
  148. ts.done(1)
  149. self.assertEqual(ts.get_ready(), ())
  150. self.assertFalse(ts.is_active())
  151. def test_is_active(self):
  152. ts = graphlib.TopologicalSorter()
  153. ts.add(1, 2)
  154. ts.prepare()
  155. self.assertTrue(ts.is_active())
  156. self.assertEqual(ts.get_ready(), (2,))
  157. self.assertTrue(ts.is_active())
  158. ts.done(2)
  159. self.assertTrue(ts.is_active())
  160. self.assertEqual(ts.get_ready(), (1,))
  161. self.assertTrue(ts.is_active())
  162. ts.done(1)
  163. self.assertFalse(ts.is_active())
  164. def test_not_hashable_nodes(self):
  165. ts = graphlib.TopologicalSorter()
  166. self.assertRaises(TypeError, ts.add, dict(), 1)
  167. self.assertRaises(TypeError, ts.add, 1, dict())
  168. self.assertRaises(TypeError, ts.add, dict(), dict())
  169. def test_order_of_insertion_does_not_matter_between_groups(self):
  170. def get_groups(ts):
  171. ts.prepare()
  172. while ts.is_active():
  173. nodes = ts.get_ready()
  174. ts.done(*nodes)
  175. yield set(nodes)
  176. ts = graphlib.TopologicalSorter()
  177. ts.add(3, 2, 1)
  178. ts.add(1, 0)
  179. ts.add(4, 5)
  180. ts.add(6, 7)
  181. ts.add(4, 7)
  182. ts2 = graphlib.TopologicalSorter()
  183. ts2.add(1, 0)
  184. ts2.add(3, 2, 1)
  185. ts2.add(4, 7)
  186. ts2.add(6, 7)
  187. ts2.add(4, 5)
  188. self.assertEqual(list(get_groups(ts)), list(get_groups(ts2)))
  189. def test_static_order_does_not_change_with_the_hash_seed(self):
  190. def check_order_with_hash_seed(seed):
  191. code = """if 1:
  192. import graphlib
  193. ts = graphlib.TopologicalSorter()
  194. ts.add('blech', 'bluch', 'hola')
  195. ts.add('abcd', 'blech', 'bluch', 'a', 'b')
  196. ts.add('a', 'a string', 'something', 'b')
  197. ts.add('bluch', 'hola', 'abcde', 'a', 'b')
  198. print(list(ts.static_order()))
  199. """
  200. env = os.environ.copy()
  201. # signal to assert_python not to do a copy
  202. # of os.environ on its own
  203. env["__cleanenv"] = True
  204. env["PYTHONHASHSEED"] = str(seed)
  205. out = assert_python_ok("-c", code, **env)
  206. return out
  207. run1 = check_order_with_hash_seed(1234)
  208. run2 = check_order_with_hash_seed(31415)
  209. self.assertNotEqual(run1, "")
  210. self.assertNotEqual(run2, "")
  211. self.assertEqual(run1, run2)
  212. if __name__ == "__main__":
  213. unittest.main()