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