@@ -457,3 +457,202 @@ class format: pass
457457 "cannot access free variable 'format' where it is not associated with a value in enclosing scope" ,
458458 ):
459459 ns ["f" ].__annotations__
460+
461+
462+ class ConditionalAnnotationTests (unittest .TestCase ):
463+ def check_scopes (self , code , true_annos , false_annos ):
464+ for scope in ("class" , "module" ):
465+ for (cond , expected ) in (
466+ # Constants (so code might get optimized out)
467+ (True , true_annos ), (False , false_annos ),
468+ # Non-constant expressions
469+ ("not not len" , true_annos ), ("not len" , false_annos ),
470+ ):
471+ with self .subTest (scope = scope , cond = cond ):
472+ code_to_run = code .format (cond = cond )
473+ if scope == "class" :
474+ code_to_run = "class Cls:\n " + textwrap .indent (textwrap .dedent (code_to_run ), " " * 4 )
475+ ns = run_code (code_to_run )
476+ if scope == "class" :
477+ self .assertEqual (ns ["Cls" ].__annotations__ , expected )
478+ else :
479+ self .assertEqual (ns ["__annotate__" ](annotationlib .Format .VALUE ),
480+ expected )
481+
482+ def test_with (self ):
483+ code = """
484+ class Swallower:
485+ def __enter__(self):
486+ pass
487+
488+ def __exit__(self, *args):
489+ return True
490+
491+ with Swallower():
492+ if {cond}:
493+ about_to_raise: int
494+ raise Exception
495+ in_with: "with"
496+ """
497+ self .check_scopes (code , {"about_to_raise" : int }, {"in_with" : "with" })
498+
499+ def test_simple_if (self ):
500+ code = """
501+ if {cond}:
502+ in_if: "if"
503+ else:
504+ in_if: "else"
505+ """
506+ self .check_scopes (code , {"in_if" : "if" }, {"in_if" : "else" })
507+
508+ def test_if_elif (self ):
509+ code = """
510+ if not len:
511+ in_if: "if"
512+ elif {cond}:
513+ in_elif: "elif"
514+ else:
515+ in_else: "else"
516+ """
517+ self .check_scopes (
518+ code ,
519+ {"in_elif" : "elif" },
520+ {"in_else" : "else" }
521+ )
522+
523+ def test_try (self ):
524+ code = """
525+ try:
526+ if {cond}:
527+ raise Exception
528+ in_try: "try"
529+ except Exception:
530+ in_except: "except"
531+ finally:
532+ in_finally: "finally"
533+ """
534+ self .check_scopes (
535+ code ,
536+ {"in_except" : "except" , "in_finally" : "finally" },
537+ {"in_try" : "try" , "in_finally" : "finally" }
538+ )
539+
540+ def test_try_star (self ):
541+ code = """
542+ try:
543+ if {cond}:
544+ raise Exception
545+ in_try_star: "try"
546+ except* Exception:
547+ in_except_star: "except"
548+ finally:
549+ in_finally: "finally"
550+ """
551+ self .check_scopes (
552+ code ,
553+ {"in_except_star" : "except" , "in_finally" : "finally" },
554+ {"in_try_star" : "try" , "in_finally" : "finally" }
555+ )
556+
557+ def test_while (self ):
558+ code = """
559+ while {cond}:
560+ in_while: "while"
561+ break
562+ else:
563+ in_else: "else"
564+ """
565+ self .check_scopes (
566+ code ,
567+ {"in_while" : "while" },
568+ {"in_else" : "else" }
569+ )
570+
571+ def test_for (self ):
572+ code = """
573+ for _ in ([1] if {cond} else []):
574+ in_for: "for"
575+ else:
576+ in_else: "else"
577+ """
578+ self .check_scopes (
579+ code ,
580+ {"in_for" : "for" , "in_else" : "else" },
581+ {"in_else" : "else" }
582+ )
583+
584+ def test_match (self ):
585+ code = """
586+ match {cond}:
587+ case True:
588+ x: "true"
589+ case False:
590+ x: "false"
591+ """
592+ self .check_scopes (
593+ code ,
594+ {"x" : "true" },
595+ {"x" : "false" }
596+ )
597+
598+ def test_nesting_override (self ):
599+ code = """
600+ if {cond}:
601+ x: "foo"
602+ if {cond}:
603+ x: "bar"
604+ """
605+ self .check_scopes (
606+ code ,
607+ {"x" : "bar" },
608+ {}
609+ )
610+
611+ def test_nesting_outer (self ):
612+ code = """
613+ if {cond}:
614+ outer_before: "outer_before"
615+ if len:
616+ inner_if: "inner_if"
617+ else:
618+ inner_else: "inner_else"
619+ outer_after: "outer_after"
620+ """
621+ self .check_scopes (
622+ code ,
623+ {"outer_before" : "outer_before" , "inner_if" : "inner_if" ,
624+ "outer_after" : "outer_after" },
625+ {}
626+ )
627+
628+ def test_nesting_inner (self ):
629+ code = """
630+ if len:
631+ outer_before: "outer_before"
632+ if {cond}:
633+ inner_if: "inner_if"
634+ else:
635+ inner_else: "inner_else"
636+ outer_after: "outer_after"
637+ """
638+ self .check_scopes (
639+ code ,
640+ {"outer_before" : "outer_before" , "inner_if" : "inner_if" ,
641+ "outer_after" : "outer_after" },
642+ {"outer_before" : "outer_before" , "inner_else" : "inner_else" ,
643+ "outer_after" : "outer_after" },
644+ )
645+
646+ def test_non_name_annotations (self ):
647+ code = """
648+ before: "before"
649+ if {cond}:
650+ a = "x"
651+ a[0]: int
652+ else:
653+ a = object()
654+ a.b: str
655+ after: "after"
656+ """
657+ expected = {"before" : "before" , "after" : "after" }
658+ self .check_scopes (code , expected , expected )
0 commit comments