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