1import unittest
2
3# For scope testing.
4g = "Global variable"
5
6
7class DictComprehensionTest(unittest.TestCase):
8
9    def test_basics(self):
10        expected = {0: 10, 1: 11, 2: 12, 3: 13, 4: 14, 5: 15, 6: 16, 7: 17,
11                    8: 18, 9: 19}
12        actual = {k: k + 10 for k in range(10)}
13        self.assertEqual(actual, expected)
14
15        expected = {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8, 9: 9}
16        actual = {k: v for k in range(10) for v in range(10) if k == v}
17        self.assertEqual(actual, expected)
18
19    def test_scope_isolation(self):
20        k = "Local Variable"
21
22        expected = {0: None, 1: None, 2: None, 3: None, 4: None, 5: None,
23                    6: None, 7: None, 8: None, 9: None}
24        actual = {k: None for k in range(10)}
25        self.assertEqual(actual, expected)
26        self.assertEqual(k, "Local Variable")
27
28        expected = {9: 1, 18: 2, 19: 2, 27: 3, 28: 3, 29: 3, 36: 4, 37: 4,
29                    38: 4, 39: 4, 45: 5, 46: 5, 47: 5, 48: 5, 49: 5, 54: 6,
30                    55: 6, 56: 6, 57: 6, 58: 6, 59: 6, 63: 7, 64: 7, 65: 7,
31                    66: 7, 67: 7, 68: 7, 69: 7, 72: 8, 73: 8, 74: 8, 75: 8,
32                    76: 8, 77: 8, 78: 8, 79: 8, 81: 9, 82: 9, 83: 9, 84: 9,
33                    85: 9, 86: 9, 87: 9, 88: 9, 89: 9}
34        actual = {k: v for v in range(10) for k in range(v * 9, v * 10)}
35        self.assertEqual(k, "Local Variable")
36        self.assertEqual(actual, expected)
37
38    def test_scope_isolation_from_global(self):
39        expected = {0: None, 1: None, 2: None, 3: None, 4: None, 5: None,
40                    6: None, 7: None, 8: None, 9: None}
41        actual = {g: None for g in range(10)}
42        self.assertEqual(actual, expected)
43        self.assertEqual(g, "Global variable")
44
45        expected = {9: 1, 18: 2, 19: 2, 27: 3, 28: 3, 29: 3, 36: 4, 37: 4,
46                    38: 4, 39: 4, 45: 5, 46: 5, 47: 5, 48: 5, 49: 5, 54: 6,
47                    55: 6, 56: 6, 57: 6, 58: 6, 59: 6, 63: 7, 64: 7, 65: 7,
48                    66: 7, 67: 7, 68: 7, 69: 7, 72: 8, 73: 8, 74: 8, 75: 8,
49                    76: 8, 77: 8, 78: 8, 79: 8, 81: 9, 82: 9, 83: 9, 84: 9,
50                    85: 9, 86: 9, 87: 9, 88: 9, 89: 9}
51        actual = {g: v for v in range(10) for g in range(v * 9, v * 10)}
52        self.assertEqual(g, "Global variable")
53        self.assertEqual(actual, expected)
54
55    def test_global_visibility(self):
56        expected = {0: 'Global variable', 1: 'Global variable',
57                    2: 'Global variable', 3: 'Global variable',
58                    4: 'Global variable', 5: 'Global variable',
59                    6: 'Global variable', 7: 'Global variable',
60                    8: 'Global variable', 9: 'Global variable'}
61        actual = {k: g for k in range(10)}
62        self.assertEqual(actual, expected)
63
64    def test_local_visibility(self):
65        v = "Local variable"
66        expected = {0: 'Local variable', 1: 'Local variable',
67                    2: 'Local variable', 3: 'Local variable',
68                    4: 'Local variable', 5: 'Local variable',
69                    6: 'Local variable', 7: 'Local variable',
70                    8: 'Local variable', 9: 'Local variable'}
71        actual = {k: v for k in range(10)}
72        self.assertEqual(actual, expected)
73        self.assertEqual(v, "Local variable")
74
75    def test_illegal_assignment(self):
76        with self.assertRaisesRegex(SyntaxError, "cannot assign"):
77            compile("{x: y for y, x in ((1, 2), (3, 4))} = 5", "<test>",
78                    "exec")
79
80        with self.assertRaisesRegex(SyntaxError, "illegal expression"):
81            compile("{x: y for y, x in ((1, 2), (3, 4))} += 5", "<test>",
82                    "exec")
83
84    def test_evaluation_order(self):
85        expected = {
86            'H': 'W',
87            'e': 'o',
88            'l': 'l',
89            'o': 'd',
90        }
91
92        expected_calls = [
93            ('key', 'H'), ('value', 'W'),
94            ('key', 'e'), ('value', 'o'),
95            ('key', 'l'), ('value', 'r'),
96            ('key', 'l'), ('value', 'l'),
97            ('key', 'o'), ('value', 'd'),
98        ]
99
100        actual_calls = []
101
102        def add_call(pos, value):
103            actual_calls.append((pos, value))
104            return value
105
106        actual = {
107            add_call('key', k): add_call('value', v)
108            for k, v in zip('Hello', 'World')
109        }
110
111        self.assertEqual(actual, expected)
112        self.assertEqual(actual_calls, expected_calls)
113
114    def test_assignment_idiom_in_comprehensions(self):
115        expected = {1: 1, 2: 4, 3: 9, 4: 16}
116        actual = {j: j*j for i in range(4) for j in [i+1]}
117        self.assertEqual(actual, expected)
118        expected = {3: 2, 5: 6, 7: 12, 9: 20}
119        actual = {j+k: j*k for i in range(4) for j in [i+1] for k in [j+1]}
120        self.assertEqual(actual, expected)
121        expected = {3: 2, 5: 6, 7: 12, 9: 20}
122        actual = {j+k: j*k for i in range(4)  for j, k in [(i+1, i+2)]}
123        self.assertEqual(actual, expected)
124
125    def test_star_expression(self):
126        expected = {0: 0, 1: 1, 2: 4, 3: 9}
127        self.assertEqual({i: i*i for i in [*range(4)]}, expected)
128        self.assertEqual({i: i*i for i in (*range(4),)}, expected)
129
130
131if __name__ == "__main__":
132    unittest.main()
133