• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import traceback
2import unittest
3
4from test.support import BrokenIter
5
6# For scope testing.
7g = "Global variable"
8
9
10class DictComprehensionTest(unittest.TestCase):
11
12    def test_basics(self):
13        expected = {0: 10, 1: 11, 2: 12, 3: 13, 4: 14, 5: 15, 6: 16, 7: 17,
14                    8: 18, 9: 19}
15        actual = {k: k + 10 for k in range(10)}
16        self.assertEqual(actual, expected)
17
18        expected = {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8, 9: 9}
19        actual = {k: v for k in range(10) for v in range(10) if k == v}
20        self.assertEqual(actual, expected)
21
22    def test_scope_isolation(self):
23        k = "Local Variable"
24
25        expected = {0: None, 1: None, 2: None, 3: None, 4: None, 5: None,
26                    6: None, 7: None, 8: None, 9: None}
27        actual = {k: None for k in range(10)}
28        self.assertEqual(actual, expected)
29        self.assertEqual(k, "Local Variable")
30
31        expected = {9: 1, 18: 2, 19: 2, 27: 3, 28: 3, 29: 3, 36: 4, 37: 4,
32                    38: 4, 39: 4, 45: 5, 46: 5, 47: 5, 48: 5, 49: 5, 54: 6,
33                    55: 6, 56: 6, 57: 6, 58: 6, 59: 6, 63: 7, 64: 7, 65: 7,
34                    66: 7, 67: 7, 68: 7, 69: 7, 72: 8, 73: 8, 74: 8, 75: 8,
35                    76: 8, 77: 8, 78: 8, 79: 8, 81: 9, 82: 9, 83: 9, 84: 9,
36                    85: 9, 86: 9, 87: 9, 88: 9, 89: 9}
37        actual = {k: v for v in range(10) for k in range(v * 9, v * 10)}
38        self.assertEqual(k, "Local Variable")
39        self.assertEqual(actual, expected)
40
41    def test_scope_isolation_from_global(self):
42        expected = {0: None, 1: None, 2: None, 3: None, 4: None, 5: None,
43                    6: None, 7: None, 8: None, 9: None}
44        actual = {g: None for g in range(10)}
45        self.assertEqual(actual, expected)
46        self.assertEqual(g, "Global variable")
47
48        expected = {9: 1, 18: 2, 19: 2, 27: 3, 28: 3, 29: 3, 36: 4, 37: 4,
49                    38: 4, 39: 4, 45: 5, 46: 5, 47: 5, 48: 5, 49: 5, 54: 6,
50                    55: 6, 56: 6, 57: 6, 58: 6, 59: 6, 63: 7, 64: 7, 65: 7,
51                    66: 7, 67: 7, 68: 7, 69: 7, 72: 8, 73: 8, 74: 8, 75: 8,
52                    76: 8, 77: 8, 78: 8, 79: 8, 81: 9, 82: 9, 83: 9, 84: 9,
53                    85: 9, 86: 9, 87: 9, 88: 9, 89: 9}
54        actual = {g: v for v in range(10) for g in range(v * 9, v * 10)}
55        self.assertEqual(g, "Global variable")
56        self.assertEqual(actual, expected)
57
58    def test_global_visibility(self):
59        expected = {0: 'Global variable', 1: 'Global variable',
60                    2: 'Global variable', 3: 'Global variable',
61                    4: 'Global variable', 5: 'Global variable',
62                    6: 'Global variable', 7: 'Global variable',
63                    8: 'Global variable', 9: 'Global variable'}
64        actual = {k: g for k in range(10)}
65        self.assertEqual(actual, expected)
66
67    def test_local_visibility(self):
68        v = "Local variable"
69        expected = {0: 'Local variable', 1: 'Local variable',
70                    2: 'Local variable', 3: 'Local variable',
71                    4: 'Local variable', 5: 'Local variable',
72                    6: 'Local variable', 7: 'Local variable',
73                    8: 'Local variable', 9: 'Local variable'}
74        actual = {k: v for k in range(10)}
75        self.assertEqual(actual, expected)
76        self.assertEqual(v, "Local variable")
77
78    def test_illegal_assignment(self):
79        with self.assertRaisesRegex(SyntaxError, "cannot assign"):
80            compile("{x: y for y, x in ((1, 2), (3, 4))} = 5", "<test>",
81                    "exec")
82
83        with self.assertRaisesRegex(SyntaxError, "illegal expression"):
84            compile("{x: y for y, x in ((1, 2), (3, 4))} += 5", "<test>",
85                    "exec")
86
87    def test_evaluation_order(self):
88        expected = {
89            'H': 'W',
90            'e': 'o',
91            'l': 'l',
92            'o': 'd',
93        }
94
95        expected_calls = [
96            ('key', 'H'), ('value', 'W'),
97            ('key', 'e'), ('value', 'o'),
98            ('key', 'l'), ('value', 'r'),
99            ('key', 'l'), ('value', 'l'),
100            ('key', 'o'), ('value', 'd'),
101        ]
102
103        actual_calls = []
104
105        def add_call(pos, value):
106            actual_calls.append((pos, value))
107            return value
108
109        actual = {
110            add_call('key', k): add_call('value', v)
111            for k, v in zip('Hello', 'World')
112        }
113
114        self.assertEqual(actual, expected)
115        self.assertEqual(actual_calls, expected_calls)
116
117    def test_assignment_idiom_in_comprehensions(self):
118        expected = {1: 1, 2: 4, 3: 9, 4: 16}
119        actual = {j: j*j for i in range(4) for j in [i+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 in [i+1] for k in [j+1]}
123        self.assertEqual(actual, expected)
124        expected = {3: 2, 5: 6, 7: 12, 9: 20}
125        actual = {j+k: j*k for i in range(4)  for j, k in [(i+1, i+2)]}
126        self.assertEqual(actual, expected)
127
128    def test_star_expression(self):
129        expected = {0: 0, 1: 1, 2: 4, 3: 9}
130        self.assertEqual({i: i*i for i in [*range(4)]}, expected)
131        self.assertEqual({i: i*i for i in (*range(4),)}, expected)
132
133    def test_exception_locations(self):
134        # The location of an exception raised from __init__ or
135        # __next__ should should be the iterator expression
136        def init_raises():
137            try:
138                {x:x for x in BrokenIter(init_raises=True)}
139            except Exception as e:
140                return e
141
142        def next_raises():
143            try:
144                {x:x for x in BrokenIter(next_raises=True)}
145            except Exception as e:
146                return e
147
148        def iter_raises():
149            try:
150                {x:x for x in BrokenIter(iter_raises=True)}
151            except Exception as e:
152                return e
153
154        for func, expected in [(init_raises, "BrokenIter(init_raises=True)"),
155                               (next_raises, "BrokenIter(next_raises=True)"),
156                               (iter_raises, "BrokenIter(iter_raises=True)"),
157                              ]:
158            with self.subTest(func):
159                exc = func()
160                f = traceback.extract_tb(exc.__traceback__)[0]
161                indent = 16
162                co = func.__code__
163                self.assertEqual(f.lineno, co.co_firstlineno + 2)
164                self.assertEqual(f.end_lineno, co.co_firstlineno + 2)
165                self.assertEqual(f.line[f.colno - indent : f.end_colno - indent],
166                                 expected)
167
168
169if __name__ == "__main__":
170    unittest.main()
171