Fix conditional execution of class and module scoped annotations

Author: JelleZijlstra

Include/internal/pycore_compile.h

@@ -122,7 +122,9 @@ void _PyCompile_ExitScope(struct _PyCompiler *c);
 Py_ssize_t _PyCompile_AddConst(struct _PyCompiler *c, PyObject *o);
 _PyInstructionSequence *_PyCompile_InstrSequence(struct _PyCompiler *c);
 int _PyCompile_FutureFeatures(struct _PyCompiler *c);
-PyObject *_PyCompile_DeferredAnnotations(struct _PyCompiler *c);
+void _PyCompile_DeferredAnnotations(
+    struct _PyCompiler *c, PyObject **deferred_annotations,
+    PyObject **conditional_annotation_indices);
 PyObject *_PyCompile_Mangle(struct _PyCompiler *c, PyObject *name);
 PyObject *_PyCompile_MaybeMangle(struct _PyCompiler *c, PyObject *name);
 int _PyCompile_MaybeAddStaticAttributeToClass(struct _PyCompiler *c, expr_ty e);
@@ -166,13 +168,16 @@ int _PyCompile_TweakInlinedComprehensionScopes(struct _PyCompiler *c, _Py_Source
                                                _PyCompile_InlinedComprehensionState *state);
 int _PyCompile_RevertInlinedComprehensionScopes(struct _PyCompiler *c, _Py_SourceLocation loc,
                                                 _PyCompile_InlinedComprehensionState *state);
-int _PyCompile_AddDeferredAnnotaion(struct _PyCompiler *c, stmt_ty s);
+int _PyCompile_AddDeferredAnnotation(struct _PyCompiler *c, stmt_ty s,
+                                     PyObject **conditional_annotation_index);
+void _PyCompile_EnterConditionalBlock(struct _PyCompiler *c);
+void _PyCompile_LeaveConditionalBlock(struct _PyCompiler *c);
 
 int _PyCodegen_AddReturnAtEnd(struct _PyCompiler *c, int addNone);
 int _PyCodegen_EnterAnonymousScope(struct _PyCompiler* c, mod_ty mod);
 int _PyCodegen_Expression(struct _PyCompiler *c, expr_ty e);
-int _PyCodegen_Body(struct _PyCompiler *c, _Py_SourceLocation loc, asdl_stmt_seq *stmts,
-                    bool is_interactive);
+int _PyCodegen_Module(struct _PyCompiler *c, _Py_SourceLocation loc, asdl_stmt_seq *stmts,
+                      bool is_interactive);
 
 /* Utility for a number of growing arrays used in the compiler */
 int _PyCompile_EnsureArrayLargeEnough(

Include/internal/pycore_global_objects_fini_generated.h

@@ -93,6 +93,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(__classdict__)
         STRUCT_FOR_ID(__classdictcell__)
         STRUCT_FOR_ID(__complex__)
+        STRUCT_FOR_ID(__conditional_annotations__)
         STRUCT_FOR_ID(__contains__)
         STRUCT_FOR_ID(__ctypes_from_outparam__)
         STRUCT_FOR_ID(__del__)

Include/internal/pycore_global_strings.h

@@ -123,6 +123,8 @@ typedef struct _symtable_entry {
                                              enclosing class scope */
     unsigned ste_has_docstring : 1; /* true if docstring present */
     unsigned ste_method : 1; /* true if block is a function block defined in class scope */
+    unsigned ste_has_conditional_annotations : 1; /* true if block has conditionally executed annotations */
+    unsigned ste_in_conditional_block : 1; /* set while we are inside a conditionally executed block */
     int ste_comp_iter_expr; /* non-zero if visiting a comprehension range expression */
     _Py_SourceLocation ste_loc; /* source location of block */
     struct _symtable_entry *ste_annotation_block; /* symbol table entry for this entry's annotations */

Include/internal/pycore_runtime_init_generated.h

@@ -457,3 +457,142 @@ class format: pass
             "cannot access free variable 'format' where it is not associated with a value in enclosing scope",
         ):
             ns["f"].__annotations__
+
+
+class ConditionalAnnotationTests(unittest.TestCase):
+    def check_scopes(self, code, true_annos, false_annos):
+        for scope in ("class", "module"):
+            for (cond, expected) in ((True, true_annos), (False, false_annos)):
+                with self.subTest(scope=scope, cond=cond):
+                    code_to_run = code.format(cond=cond)
+                    if scope == "class":
+                        code_to_run = "class Cls:\n" + textwrap.indent(textwrap.dedent(code_to_run), " " * 4)
+                    ns = run_code(code_to_run)
+                    if scope == "class":
+                        self.assertEqual(ns["Cls"].__annotations__, expected)
+                    else:
+                        self.assertEqual(ns["__annotate__"](annotationlib.Format.VALUE),
+                                         expected)
+
+    def test_with(self):
+        code = """
+            import contextlib
+            class Swallower:
+                def __enter__(self):
+                    pass
+
+                def __exit__(self, *args):
+                    return True
+
+            with Swallower():
+                if {cond}:
+                    about_to_raise: int
+                    raise Exception
+                in_with: "with"
+        """
+        self.check_scopes(code, {"about_to_raise": int}, {"in_with": "with"})
+
+    def test_simple_if(self):
+        code = """
+            if {cond}:
+                in_if: "if"
+            else:
+                in_if: "else"
+        """
+        self.check_scopes(code, {"in_if": "if"}, {"in_if": "else"})
+
+    def test_try(self):
+        code = """
+            try:
+                if {cond}:
+                    raise Exception
+                in_try: "try"
+            except Exception:
+                in_except: "except"
+            finally:
+                in_finally: "finally"
+        """
+        self.check_scopes(
+            code,
+            {"in_except": "except", "in_finally": "finally"},
+            {"in_try": "try", "in_finally": "finally"}
+        )
+
+    def test_try_star(self):
+        code = """
+            try:
+                if {cond}:
+                    raise Exception
+                in_try_star: "try"
+            except* Exception:
+                in_except_star: "except"
+            finally:
+                in_finally: "finally"
+        """
+        self.check_scopes(
+            code,
+            {"in_except_star": "except", "in_finally": "finally"},
+            {"in_try_star": "try", "in_finally": "finally"}
+        )
+
+    def test_while(self):
+        code = """
+            while {cond}:
+                in_while: "while"
+                break
+            else:
+                in_else: "else"
+        """
+        self.check_scopes(
+            code,
+            {"in_while": "while"},
+            {"in_else": "else"}
+        )
+
+    def test_for(self):
+        code = """
+            for _ in ([1] if {cond} else []):
+                in_for: "for"
+            else:
+                in_else: "else"
+        """
+        self.check_scopes(
+            code,
+            {"in_for": "for", "in_else": "else"},
+            {"in_else": "else"}
+        )
+
+    def test_nesting_outer(self):
+        code = """
+            if {cond}:
+                outer_before: "outer_before"
+                if len:
+                    inner_if: "inner_if"
+                else:
+                    inner_else: "inner_else"
+                outer_after: "outer_after"
+        """
+        self.check_scopes(
+            code,
+            {"outer_before": "outer_before", "inner_if": "inner_if",
+             "outer_after": "outer_after"},
+            {}
+        )
+
+    def test_nesting_inner(self):
+        code = """
+            if len:
+                outer_before: "outer_before"
+                if {cond}:
+                    inner_if: "inner_if"
+                else:
+                    inner_else: "inner_else"
+                outer_after: "outer_after"
+        """
+        self.check_scopes(
+            code,
+            {"outer_before": "outer_before", "inner_if": "inner_if",
+             "outer_after": "outer_after"},
+            {"outer_before": "outer_before", "inner_else": "inner_else",
+             "outer_after": "outer_after"},
+        )