1 import doctest
2 import textwrap
3 import unittest
4
5
6 doctests = """
7 ########### Tests borrowed from or inspired by test_genexps.py ############
8
9 Test simple loop with conditional
10
11 >>> sum([i*i for i in range(100) if i&1 == 1])
12 166650
13
14 Test simple nesting
15
16 >>> [(i,j) for i in range(3) for j in range(4)]
17 [(0, 0), (0, 1), (0, 2), (0, 3), (1, 0), (1, 1), (1, 2), (1, 3), (2, 0), (2, 1), (2, 2), (2, 3)]
18
19 Test nesting with the inner expression dependent on the outer
20
21 >>> [(i,j) for i in range(4) for j in range(i)]
22 [(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2)]
23
24 Test the idiom for temporary variable assignment in comprehensions.
25
26 >>> [j*j for i in range(4) for j in [i+1]]
27 [1, 4, 9, 16]
28 >>> [j*k for i in range(4) for j in [i+1] for k in [j+1]]
29 [2, 6, 12, 20]
30 >>> [j*k for i in range(4) for j, k in [(i+1, i+2)]]
31 [2, 6, 12, 20]
32
33 Not assignment
34
35 >>> [i*i for i in [*range(4)]]
36 [0, 1, 4, 9]
37 >>> [i*i for i in (*range(4),)]
38 [0, 1, 4, 9]
39
40 Make sure the induction variable is not exposed
41
42 >>> i = 20
43 >>> sum([i*i for i in range(100)])
44 328350
45
46 >>> i
47 20
48
49 Verify that syntax error's are raised for listcomps used as lvalues
50
51 >>> [y for y in (1,2)] = 10 # doctest: +IGNORE_EXCEPTION_DETAIL
52 Traceback (most recent call last):
53 ...
54 SyntaxError: ...
55
56 >>> [y for y in (1,2)] += 10 # doctest: +IGNORE_EXCEPTION_DETAIL
57 Traceback (most recent call last):
58 ...
59 SyntaxError: ...
60
61
62 ########### Tests borrowed from or inspired by test_generators.py ############
63
64 Make a nested list comprehension that acts like range()
65
66 >>> def frange(n):
67 ... return [i for i in range(n)]
68 >>> frange(10)
69 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
70
71 Same again, only as a lambda expression instead of a function definition
72
73 >>> lrange = lambda n: [i for i in range(n)]
74 >>> lrange(10)
75 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
76
77 Generators can call other generators:
78
79 >>> def grange(n):
80 ... for x in [i for i in range(n)]:
81 ... yield x
82 >>> list(grange(5))
83 [0, 1, 2, 3, 4]
84
85
86 Make sure that None is a valid return value
87
88 >>> [None for i in range(10)]
89 [None, None, None, None, None, None, None, None, None, None]
90
91 """
92
93
94 class ESC[4;38;5;81mListComprehensionTest(ESC[4;38;5;149munittestESC[4;38;5;149m.ESC[4;38;5;149mTestCase):
95 def _check_in_scopes(self, code, outputs=None, ns=None, scopes=None, raises=()):
96 code = textwrap.dedent(code)
97 scopes = scopes or ["module", "class", "function"]
98 for scope in scopes:
99 with self.subTest(scope=scope):
100 if scope == "class":
101 newcode = textwrap.dedent("""
102 class _C:
103 {code}
104 """).format(code=textwrap.indent(code, " "))
105 def get_output(moddict, name):
106 return getattr(moddict["_C"], name)
107 elif scope == "function":
108 newcode = textwrap.dedent("""
109 def _f():
110 {code}
111 return locals()
112 _out = _f()
113 """).format(code=textwrap.indent(code, " "))
114 def get_output(moddict, name):
115 return moddict["_out"][name]
116 else:
117 newcode = code
118 def get_output(moddict, name):
119 return moddict[name]
120 newns = ns.copy() if ns else {}
121 try:
122 exec(newcode, newns)
123 except raises as e:
124 # We care about e.g. NameError vs UnboundLocalError
125 self.assertIs(type(e), raises)
126 else:
127 for k, v in (outputs or {}).items():
128 self.assertEqual(get_output(newns, k), v, k)
129
130 def test_lambdas_with_iteration_var_as_default(self):
131 code = """
132 items = [(lambda i=i: i) for i in range(5)]
133 y = [x() for x in items]
134 """
135 outputs = {"y": [0, 1, 2, 3, 4]}
136 self._check_in_scopes(code, outputs)
137
138 def test_lambdas_with_free_var(self):
139 code = """
140 items = [(lambda: i) for i in range(5)]
141 y = [x() for x in items]
142 """
143 outputs = {"y": [4, 4, 4, 4, 4]}
144 self._check_in_scopes(code, outputs)
145
146 def test_class_scope_free_var_with_class_cell(self):
147 class ESC[4;38;5;81mC:
148 def method(self):
149 super()
150 return __class__
151 items = [(lambda: i) for i in range(5)]
152 y = [x() for x in items]
153
154 self.assertEqual(C.y, [4, 4, 4, 4, 4])
155 self.assertIs(C().method(), C)
156
157 def test_inner_cell_shadows_outer(self):
158 code = """
159 items = [(lambda: i) for i in range(5)]
160 i = 20
161 y = [x() for x in items]
162 """
163 outputs = {"y": [4, 4, 4, 4, 4], "i": 20}
164 self._check_in_scopes(code, outputs)
165
166 def test_inner_cell_shadows_outer_no_store(self):
167 code = """
168 def f(x):
169 return [lambda: x for x in range(x)], x
170 fns, x = f(2)
171 y = [fn() for fn in fns]
172 """
173 outputs = {"y": [1, 1], "x": 2}
174 self._check_in_scopes(code, outputs)
175
176 def test_closure_can_jump_over_comp_scope(self):
177 code = """
178 items = [(lambda: y) for i in range(5)]
179 y = 2
180 z = [x() for x in items]
181 """
182 outputs = {"z": [2, 2, 2, 2, 2]}
183 self._check_in_scopes(code, outputs, scopes=["module", "function"])
184
185 def test_cell_inner_free_outer(self):
186 code = """
187 def f():
188 return [lambda: x for x in (x, [1])[1]]
189 x = ...
190 y = [fn() for fn in f()]
191 """
192 outputs = {"y": [1]}
193 self._check_in_scopes(code, outputs, scopes=["module", "function"])
194
195 def test_free_inner_cell_outer(self):
196 code = """
197 g = 2
198 def f():
199 return g
200 y = [g for x in [1]]
201 """
202 outputs = {"y": [2]}
203 self._check_in_scopes(code, outputs, scopes=["module", "function"])
204 self._check_in_scopes(code, scopes=["class"], raises=NameError)
205
206 def test_inner_cell_shadows_outer_redefined(self):
207 code = """
208 y = 10
209 items = [(lambda: y) for y in range(5)]
210 x = y
211 y = 20
212 out = [z() for z in items]
213 """
214 outputs = {"x": 10, "out": [4, 4, 4, 4, 4]}
215 self._check_in_scopes(code, outputs)
216
217 def test_shadows_outer_cell(self):
218 code = """
219 def inner():
220 return g
221 [g for g in range(5)]
222 x = inner()
223 """
224 outputs = {"x": -1}
225 self._check_in_scopes(code, outputs, ns={"g": -1})
226
227 def test_explicit_global(self):
228 code = """
229 global g
230 x = g
231 g = 2
232 items = [g for g in [1]]
233 y = g
234 """
235 outputs = {"x": 1, "y": 2, "items": [1]}
236 self._check_in_scopes(code, outputs, ns={"g": 1})
237
238 def test_explicit_global_2(self):
239 code = """
240 global g
241 x = g
242 g = 2
243 items = [g for x in [1]]
244 y = g
245 """
246 outputs = {"x": 1, "y": 2, "items": [2]}
247 self._check_in_scopes(code, outputs, ns={"g": 1})
248
249 def test_explicit_global_3(self):
250 code = """
251 global g
252 fns = [lambda: g for g in [2]]
253 items = [fn() for fn in fns]
254 """
255 outputs = {"items": [2]}
256 self._check_in_scopes(code, outputs, ns={"g": 1})
257
258 def test_assignment_expression(self):
259 code = """
260 x = -1
261 items = [(x:=y) for y in range(3)]
262 """
263 outputs = {"x": 2}
264 # assignment expression in comprehension is disallowed in class scope
265 self._check_in_scopes(code, outputs, scopes=["module", "function"])
266
267 def test_free_var_in_comp_child(self):
268 code = """
269 lst = range(3)
270 funcs = [lambda: x for x in lst]
271 inc = [x + 1 for x in lst]
272 [x for x in inc]
273 x = funcs[0]()
274 """
275 outputs = {"x": 2}
276 self._check_in_scopes(code, outputs)
277
278 def test_shadow_with_free_and_local(self):
279 code = """
280 lst = range(3)
281 x = -1
282 funcs = [lambda: x for x in lst]
283 items = [x + 1 for x in lst]
284 """
285 outputs = {"x": -1}
286 self._check_in_scopes(code, outputs)
287
288 def test_shadow_comp_iterable_name(self):
289 code = """
290 x = [1]
291 y = [x for x in x]
292 """
293 outputs = {"x": [1]}
294 self._check_in_scopes(code, outputs)
295
296 def test_nested_free(self):
297 code = """
298 x = 1
299 def g():
300 [x for x in range(3)]
301 return x
302 g()
303 """
304 outputs = {"x": 1}
305 self._check_in_scopes(code, outputs, scopes=["module", "function"])
306
307 def test_introspecting_frame_locals(self):
308 code = """
309 import sys
310 [i for i in range(2)]
311 i = 20
312 sys._getframe().f_locals
313 """
314 outputs = {"i": 20}
315 self._check_in_scopes(code, outputs)
316
317 def test_nested(self):
318 code = """
319 l = [2, 3]
320 y = [[x ** 2 for x in range(x)] for x in l]
321 """
322 outputs = {"y": [[0, 1], [0, 1, 4]]}
323 self._check_in_scopes(code, outputs)
324
325 def test_nested_2(self):
326 code = """
327 l = [1, 2, 3]
328 x = 3
329 y = [x for [x ** x for x in range(x)][x - 1] in l]
330 """
331 outputs = {"y": [3, 3, 3]}
332 self._check_in_scopes(code, outputs, scopes=["module", "function"])
333 self._check_in_scopes(code, scopes=["class"], raises=NameError)
334
335 def test_nested_3(self):
336 code = """
337 l = [(1, 2), (3, 4), (5, 6)]
338 y = [x for (x, [x ** x for x in range(x)][x - 1]) in l]
339 """
340 outputs = {"y": [1, 3, 5]}
341 self._check_in_scopes(code, outputs)
342
343 def test_nested_4(self):
344 code = """
345 items = [([lambda: x for x in range(2)], lambda: x) for x in range(3)]
346 out = [([fn() for fn in fns], fn()) for fns, fn in items]
347 """
348 outputs = {"out": [([1, 1], 2), ([1, 1], 2), ([1, 1], 2)]}
349 self._check_in_scopes(code, outputs)
350
351 def test_nameerror(self):
352 code = """
353 [x for x in [1]]
354 x
355 """
356
357 self._check_in_scopes(code, raises=NameError)
358
359 def test_dunder_name(self):
360 code = """
361 y = [__x for __x in [1]]
362 """
363 outputs = {"y": [1]}
364 self._check_in_scopes(code, outputs)
365
366 def test_unbound_local_after_comprehension(self):
367 def f():
368 if False:
369 x = 0
370 [x for x in [1]]
371 return x
372
373 with self.assertRaises(UnboundLocalError):
374 f()
375
376 def test_unbound_local_inside_comprehension(self):
377 def f():
378 l = [None]
379 return [1 for (l[0], l) in [[1, 2]]]
380
381 with self.assertRaises(UnboundLocalError):
382 f()
383
384 def test_global_outside_cellvar_inside_plus_freevar(self):
385 code = """
386 a = 1
387 def f():
388 func, = [(lambda: b) for b in [a]]
389 return b, func()
390 x = f()
391 """
392 self._check_in_scopes(
393 code, {"x": (2, 1)}, ns={"b": 2}, scopes=["function", "module"])
394 # inside a class, the `a = 1` assignment is not visible
395 self._check_in_scopes(code, raises=NameError, scopes=["class"])
396
397 def test_cell_in_nested_comprehension(self):
398 code = """
399 a = 1
400 def f():
401 (func, inner_b), = [[lambda: b for b in c] + [b] for c in [[a]]]
402 return b, inner_b, func()
403 x = f()
404 """
405 self._check_in_scopes(
406 code, {"x": (2, 2, 1)}, ns={"b": 2}, scopes=["function", "module"])
407 # inside a class, the `a = 1` assignment is not visible
408 self._check_in_scopes(code, raises=NameError, scopes=["class"])
409
410 def test_name_error_in_class_scope(self):
411 code = """
412 y = 1
413 [x + y for x in range(2)]
414 """
415 self._check_in_scopes(code, raises=NameError, scopes=["class"])
416
417 def test_global_in_class_scope(self):
418 code = """
419 y = 2
420 vals = [(x, y) for x in range(2)]
421 """
422 outputs = {"vals": [(0, 1), (1, 1)]}
423 self._check_in_scopes(code, outputs, ns={"y": 1}, scopes=["class"])
424
425 def test_in_class_scope_inside_function_1(self):
426 code = """
427 class C:
428 y = 2
429 vals = [(x, y) for x in range(2)]
430 vals = C.vals
431 """
432 outputs = {"vals": [(0, 1), (1, 1)]}
433 self._check_in_scopes(code, outputs, ns={"y": 1}, scopes=["function"])
434
435 def test_in_class_scope_inside_function_2(self):
436 code = """
437 y = 1
438 class C:
439 y = 2
440 vals = [(x, y) for x in range(2)]
441 vals = C.vals
442 """
443 outputs = {"vals": [(0, 1), (1, 1)]}
444 self._check_in_scopes(code, outputs, scopes=["function"])
445
446 def test_in_class_scope_with_global(self):
447 code = """
448 y = 1
449 class C:
450 global y
451 y = 2
452 # Ensure the listcomp uses the global, not the value in the
453 # class namespace
454 locals()['y'] = 3
455 vals = [(x, y) for x in range(2)]
456 vals = C.vals
457 """
458 outputs = {"vals": [(0, 2), (1, 2)]}
459 self._check_in_scopes(code, outputs, scopes=["module", "class"])
460 outputs = {"vals": [(0, 1), (1, 1)]}
461 self._check_in_scopes(code, outputs, scopes=["function"])
462
463 def test_in_class_scope_with_nonlocal(self):
464 code = """
465 y = 1
466 class C:
467 nonlocal y
468 y = 2
469 # Ensure the listcomp uses the global, not the value in the
470 # class namespace
471 locals()['y'] = 3
472 vals = [(x, y) for x in range(2)]
473 vals = C.vals
474 """
475 outputs = {"vals": [(0, 2), (1, 2)]}
476 self._check_in_scopes(code, outputs, scopes=["function"])
477
478 def test_nested_has_free_var(self):
479 code = """
480 items = [a for a in [1] if [a for _ in [0]]]
481 """
482 outputs = {"items": [1]}
483 self._check_in_scopes(code, outputs, scopes=["class"])
484
485 def test_nested_free_var_not_bound_in_outer_comp(self):
486 code = """
487 z = 1
488 items = [a for a in [1] if [x for x in [1] if z]]
489 """
490 self._check_in_scopes(code, {"items": [1]}, scopes=["module", "function"])
491 self._check_in_scopes(code, {"items": []}, ns={"z": 0}, scopes=["class"])
492
493 def test_nested_free_var_in_iter(self):
494 code = """
495 items = [_C for _C in [1] for [0, 1][[x for x in [1] if _C][0]] in [2]]
496 """
497 self._check_in_scopes(code, {"items": [1]})
498
499 def test_nested_free_var_in_expr(self):
500 code = """
501 items = [(_C, [x for x in [1] if _C]) for _C in [0, 1]]
502 """
503 self._check_in_scopes(code, {"items": [(0, []), (1, [1])]})
504
505 def test_nested_listcomp_in_lambda(self):
506 code = """
507 f = [(z, lambda y: [(x, y, z) for x in [3]]) for z in [1]]
508 (z, func), = f
509 out = func(2)
510 """
511 self._check_in_scopes(code, {"z": 1, "out": [(3, 2, 1)]})
512
513 def test_lambda_in_iter(self):
514 code = """
515 (func, c), = [(a, b) for b in [1] for a in [lambda : a]]
516 d = func()
517 assert d is func
518 # must use "a" in this scope
519 e = a if False else None
520 """
521 self._check_in_scopes(code, {"c": 1, "e": None})
522
523 def test_assign_to_comp_iter_var_in_outer_function(self):
524 code = """
525 a = [1 for a in [0]]
526 """
527 self._check_in_scopes(code, {"a": [1]}, scopes=["function"])
528
529 def test_no_leakage_to_locals(self):
530 code = """
531 def b():
532 [a for b in [1] for _ in []]
533 return b, locals()
534 r, s = b()
535 x = r is b
536 y = list(s.keys())
537 """
538 self._check_in_scopes(code, {"x": True, "y": []}, scopes=["module"])
539 self._check_in_scopes(code, {"x": True, "y": ["b"]}, scopes=["function"])
540 self._check_in_scopes(code, raises=NameError, scopes=["class"])
541
542 def test_iter_var_available_in_locals(self):
543 code = """
544 l = [1, 2]
545 y = 0
546 items = [locals()["x"] for x in l]
547 items2 = [vars()["x"] for x in l]
548 items3 = [("x" in dir()) for x in l]
549 items4 = [eval("x") for x in l]
550 # x is available, and does not overwrite y
551 [exec("y = x") for x in l]
552 """
553 self._check_in_scopes(
554 code,
555 {
556 "items": [1, 2],
557 "items2": [1, 2],
558 "items3": [True, True],
559 "items4": [1, 2],
560 "y": 0
561 }
562 )
563
564 def test_comp_in_try_except(self):
565 template = """
566 value = ["ab"]
567 result = snapshot = None
568 try:
569 result = [{func}(value) for value in value]
570 except:
571 snapshot = value
572 raise
573 """
574 # No exception.
575 code = template.format(func='len')
576 self._check_in_scopes(code, {"value": ["ab"], "result": [2], "snapshot": None})
577 # Handles exception.
578 code = template.format(func='int')
579 self._check_in_scopes(code, {"value": ["ab"], "result": None, "snapshot": ["ab"]},
580 raises=ValueError)
581
582 def test_comp_in_try_finally(self):
583 template = """
584 value = ["ab"]
585 result = snapshot = None
586 try:
587 result = [{func}(value) for value in value]
588 finally:
589 snapshot = value
590 """
591 # No exception.
592 code = template.format(func='len')
593 self._check_in_scopes(code, {"value": ["ab"], "result": [2], "snapshot": ["ab"]})
594 # Handles exception.
595 code = template.format(func='int')
596 self._check_in_scopes(code, {"value": ["ab"], "result": None, "snapshot": ["ab"]},
597 raises=ValueError)
598
599 def test_exception_in_post_comp_call(self):
600 code = """
601 value = [1, None]
602 try:
603 [v for v in value].sort()
604 except:
605 pass
606 """
607 self._check_in_scopes(code, {"value": [1, None]})
608
609 def test_frame_locals(self):
610 code = """
611 val = [sys._getframe().f_locals for a in [0]][0]["a"]
612 """
613 import sys
614 self._check_in_scopes(code, {"val": 0}, ns={"sys": sys})
615
616
617 __test__ = {'doctests' : doctests}
618
619 def load_tests(loader, tests, pattern):
620 tests.addTest(doctest.DocTestSuite())
621 return tests
622
623
624 if __name__ == "__main__":
625 unittest.main()