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