gh-130881: Handle conditionally defined annotations (#130935)

Include/internal/pycore_compile.h

@@ -134,7 +134,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);
@@ -178,13 +180,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,
+int _PyCodegen_Module(struct _PyCompiler *c, _Py_SourceLocation loc, asdl_stmt_seq *stmts,
-                    bool is_interactive);
+                      bool is_interactive);
 
 int _PyCompile_ConstCacheMergeOne(PyObject *const_cache, PyObject **obj);
 

Include/internal/pycore_global_objects_fini_generated.h

@@ -604,6 +604,7 @@ _PyStaticObjects_CheckRefcnt(PyInterpreterState *interp) {
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__classdict__));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__classdictcell__));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__complex__));
+    _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__conditional_annotations__));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__contains__));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__ctypes_from_outparam__));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(__del__));

Include/internal/pycore_global_strings.h

@@ -95,6 +95,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_runtime_init_generated.h

@@ -602,6 +602,7 @@ extern "C" {
     INIT_ID(__classdict__), \
     INIT_ID(__classdictcell__), \
     INIT_ID(__complex__), \
+    INIT_ID(__conditional_annotations__), \
     INIT_ID(__contains__), \
     INIT_ID(__ctypes_from_outparam__), \
     INIT_ID(__del__), \

Include/internal/pycore_symtable.h

@@ -124,6 +124,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_unicodeobject_generated.h

@@ -168,6 +168,10 @@ _PyUnicode_InitStaticStrings(PyInterpreterState *interp) {
     _PyUnicode_InternStatic(interp, &string);
     assert(_PyUnicode_CheckConsistency(string, 1));
     assert(PyUnicode_GET_LENGTH(string) != 1);
+    string = &_Py_ID(__conditional_annotations__);
+    _PyUnicode_InternStatic(interp, &string);
+    assert(_PyUnicode_CheckConsistency(string, 1));
+    assert(PyUnicode_GET_LENGTH(string) != 1);
     string = &_Py_ID(__contains__);
     _PyUnicode_InternStatic(interp, &string);
     assert(_PyUnicode_CheckConsistency(string, 1));

Lib/test/test_type_annotations.py

@@ -457,3 +457,202 @@ class DeferredEvaluationTests(unittest.TestCase):
             "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 (
+                # Constants (so code might get optimized out)
+                (True, true_annos), (False, false_annos),
+                # Non-constant expressions
+                ("not not len", true_annos), ("not len", 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 = """
+            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_if_elif(self):
+        code = """
+            if not len:
+                in_if: "if"
+            elif {cond}:
+                in_elif: "elif"
+            else:
+                in_else: "else"
+        """
+        self.check_scopes(
+            code,
+            {"in_elif": "elif"},
+            {"in_else": "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_match(self):
+        code = """
+            match {cond}:
+                case True:
+                    x: "true"
+                case False:
+                    x: "false"
+        """
+        self.check_scopes(
+            code,
+            {"x": "true"},
+            {"x": "false"}
+        )
+
+    def test_nesting_override(self):
+        code = """
+            if {cond}:
+                x: "foo"
+                if {cond}:
+                    x: "bar"
+        """
+        self.check_scopes(
+            code,
+            {"x": "bar"},
+            {}
+        )
+
+    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"},
+        )
+
+    def test_non_name_annotations(self):
+        code = """
+            before: "before"
+            if {cond}:
+                a = "x"
+                a[0]: int
+            else:
+                a = object()
+                a.b: str
+            after: "after"
+        """
+        expected = {"before": "before", "after": "after"}
+        self.check_scopes(code, expected, expected)

Misc/NEWS.d/next/Core_and_Builtins/2025-03-06-17-04-27.gh-issue-130935.ss4lmS.rst

@@ -0,0 +1,3 @@
+Annotations at the class and module level that are conditionally defined are
+now only reflected in ``__annotations__`` if the block they are in is
+executed. Patch by Jelle Zijlstra.

Python/codegen.c

@@ -2,7 +2,7 @@
  * This file implements the compiler's code generation stage, which
  * produces a sequence of pseudo-instructions from an AST.
  *
- * The primary entry point is _PyCodegen_Body() for modules, and
+ * The primary entry point is _PyCodegen_Module() for modules, and
  * _PyCodegen_Expression() for expressions.
  *
  * CAUTION: The VISIT_* macros abort the current function when they
@@ -204,8 +204,11 @@ static int codegen_subscript(compiler *, expr_ty);
 static int codegen_slice_two_parts(compiler *, expr_ty);
 static int codegen_slice(compiler *, expr_ty);
 
-static int codegen_with(compiler *, stmt_ty, int);
+static int codegen_body(compiler *, location, asdl_stmt_seq *, bool);
-static int codegen_async_with(compiler *, stmt_ty, int);
+static int codegen_with(compiler *, stmt_ty);
+static int codegen_async_with(compiler *, stmt_ty);
+static int codegen_with_inner(compiler *, stmt_ty, int);
+static int codegen_async_with_inner(compiler *, stmt_ty, int);
 static int codegen_async_for(compiler *, stmt_ty);
 static int codegen_call_simple_kw_helper(compiler *c,
                                          location loc,
@@ -681,12 +684,13 @@ codegen_setup_annotations_scope(compiler *c, location loc,
 }
 
 static int
-codegen_leave_annotations_scope(compiler *c, location loc,
+codegen_leave_annotations_scope(compiler *c, location loc)
-                                Py_ssize_t annotations_len)
 {
-    ADDOP_I(c, loc, BUILD_MAP, annotations_len);
     ADDOP_IN_SCOPE(c, loc, RETURN_VALUE);
     PyCodeObject *co = _PyCompile_OptimizeAndAssemble(c, 1);
+    if (co == NULL) {
+        return ERROR;
+    }
 
     // We want the parameter to __annotate__ to be named "format" in the
     // signature  shown by inspect.signature(), but we need to use a
@@ -696,16 +700,19 @@ codegen_leave_annotations_scope(compiler *c, location loc,
     // co->co_localsplusnames = ("format", *co->co_localsplusnames[1:])
     const Py_ssize_t size = PyObject_Size(co->co_localsplusnames);
     if (size == -1) {
+        Py_DECREF(co);
         return ERROR;
     }
     PyObject *new_names = PyTuple_New(size);
     if (new_names == NULL) {
+        Py_DECREF(co);
         return ERROR;
     }
     PyTuple_SET_ITEM(new_names, 0, Py_NewRef(&_Py_ID(format)));
     for (int i = 1; i < size; i++) {
         PyObject *item = PyTuple_GetItem(co->co_localsplusnames, i);
         if (item == NULL) {
+            Py_DECREF(co);
             Py_DECREF(new_names);
             return ERROR;
         }
@@ -715,23 +722,75 @@ codegen_leave_annotations_scope(compiler *c, location loc,
     Py_SETREF(co->co_localsplusnames, new_names);
 
     _PyCompile_ExitScope(c);
-    if (co == NULL) {
-        return ERROR;
-    }
     int ret = codegen_make_closure(c, loc, co, 0);
     Py_DECREF(co);
     RETURN_IF_ERROR(ret);
     return SUCCESS;
 }
 
+static int
+codegen_deferred_annotations_body(compiler *c, location loc,
+    PyObject *deferred_anno, PyObject *conditional_annotation_indices, int scope_type)
+{
+    Py_ssize_t annotations_len = PyList_GET_SIZE(deferred_anno);
+
+    assert(PyList_CheckExact(conditional_annotation_indices));
+    assert(annotations_len == PyList_Size(conditional_annotation_indices));
+
+    ADDOP_I(c, loc, BUILD_MAP, 0); // stack now contains <annos>
+
+    for (Py_ssize_t i = 0; i < annotations_len; i++) {
+        PyObject *ptr = PyList_GET_ITEM(deferred_anno, i);
+        stmt_ty st = (stmt_ty)PyLong_AsVoidPtr(ptr);
+        if (st == NULL) {
+            return ERROR;
+        }
+        PyObject *mangled = _PyCompile_Mangle(c, st->v.AnnAssign.target->v.Name.id);
+        if (!mangled) {
+            return ERROR;
+        }
+        PyObject *cond_index = PyList_GET_ITEM(conditional_annotation_indices, i);
+        assert(PyLong_CheckExact(cond_index));
+        long idx = PyLong_AS_LONG(cond_index);
+        NEW_JUMP_TARGET_LABEL(c, not_set);
+
+        if (idx != -1) {
+            ADDOP_LOAD_CONST(c, LOC(st), cond_index);
+            if (scope_type == COMPILE_SCOPE_CLASS) {
+                ADDOP_NAME(
+                    c, LOC(st), LOAD_DEREF, &_Py_ID(__conditional_annotations__), freevars);
+            }
+            else {
+                ADDOP_NAME(
+                    c, LOC(st), LOAD_GLOBAL, &_Py_ID(__conditional_annotations__), names);
+            }
+
+            ADDOP_I(c, LOC(st), CONTAINS_OP, 0);
+            ADDOP_JUMP(c, LOC(st), POP_JUMP_IF_FALSE, not_set);
+        }
+
+        VISIT(c, expr, st->v.AnnAssign.annotation);
+        ADDOP_I(c, LOC(st), COPY, 2);
+        ADDOP_LOAD_CONST_NEW(c, LOC(st), mangled);
+        // stack now contains <annos> <name> <annos> <value>
+        ADDOP(c, loc, STORE_SUBSCR);
+        // stack now contains <annos>
+
+        USE_LABEL(c, not_set);
+    }
+    return SUCCESS;
+}
+
 static int
 codegen_process_deferred_annotations(compiler *c, location loc)
 {
-    PyObject *deferred_anno = _PyCompile_DeferredAnnotations(c);
+    PyObject *deferred_anno = NULL;
+    PyObject *conditional_annotation_indices = NULL;
+    _PyCompile_DeferredAnnotations(c, &deferred_anno, &conditional_annotation_indices);
     if (deferred_anno == NULL) {
+        assert(conditional_annotation_indices == NULL);
         return SUCCESS;
     }
-    Py_INCREF(deferred_anno);
 
     // It's possible that ste_annotations_block is set but
     // u_deferred_annotations is not, because the former is still
@@ -741,35 +800,28 @@ codegen_process_deferred_annotations(compiler *c, location loc)
     PySTEntryObject *ste = SYMTABLE_ENTRY(c);
     assert(ste->ste_annotation_block != NULL);
     void *key = (void *)((uintptr_t)ste->ste_id + 1);
+    int scope_type = SCOPE_TYPE(c);
     if (codegen_setup_annotations_scope(c, loc, key,
                                         ste->ste_annotation_block->ste_name) < 0) {
-        Py_DECREF(deferred_anno);
+        goto error;
-        return ERROR;
     }
-    Py_ssize_t annotations_len = PyList_Size(deferred_anno);
+    if (codegen_deferred_annotations_body(c, loc, deferred_anno,
-    for (Py_ssize_t i = 0; i < annotations_len; i++) {
+                                          conditional_annotation_indices, scope_type) < 0) {
-        PyObject *ptr = PyList_GET_ITEM(deferred_anno, i);
+        _PyCompile_ExitScope(c);
-        stmt_ty st = (stmt_ty)PyLong_AsVoidPtr(ptr);
+        goto error;
-        if (st == NULL) {
-            _PyCompile_ExitScope(c);
-            Py_DECREF(deferred_anno);
-            return ERROR;
-        }
-        PyObject *mangled = _PyCompile_Mangle(c, st->v.AnnAssign.target->v.Name.id);
-        if (!mangled) {
-            _PyCompile_ExitScope(c);
-            Py_DECREF(deferred_anno);
-            return ERROR;
-        }
-        ADDOP_LOAD_CONST_NEW(c, LOC(st), mangled);
-        VISIT(c, expr, st->v.AnnAssign.annotation);
     }
-    Py_DECREF(deferred_anno);
 
-    RETURN_IF_ERROR(codegen_leave_annotations_scope(c, loc, annotations_len));
+    Py_DECREF(deferred_anno);
+    Py_DECREF(conditional_annotation_indices);
+
+    RETURN_IF_ERROR(codegen_leave_annotations_scope(c, loc));
     RETURN_IF_ERROR(codegen_nameop(c, loc, &_Py_ID(__annotate__), Store));
 
     return SUCCESS;
+error:
+    Py_XDECREF(deferred_anno);
+    Py_XDECREF(conditional_annotation_indices);
+    return ERROR;
 }
 
 /* Compile an expression */
@@ -784,7 +836,17 @@ _PyCodegen_Expression(compiler *c, expr_ty e)
    and for annotations. */
 
 int
-_PyCodegen_Body(compiler *c, location loc, asdl_stmt_seq *stmts, bool is_interactive)
+_PyCodegen_Module(compiler *c, location loc, asdl_stmt_seq *stmts, bool is_interactive)
+{
+    if (SYMTABLE_ENTRY(c)->ste_has_conditional_annotations) {
+        ADDOP_I(c, loc, BUILD_SET, 0);
+        ADDOP_N(c, loc, STORE_NAME, &_Py_ID(__conditional_annotations__), names);
+    }
+    return codegen_body(c, loc, stmts, is_interactive);
+}
+
+int
+codegen_body(compiler *c, location loc, asdl_stmt_seq *stmts, bool is_interactive)
 {
     /* If from __future__ import annotations is active,
      * every annotated class and module should have __annotations__.
@@ -1029,8 +1091,8 @@ codegen_annotations_in_scope(compiler *c, location loc,
 }
 
 static int
-codegen_annotations(compiler *c, location loc,
+codegen_function_annotations(compiler *c, location loc,
-                    arguments_ty args, expr_ty returns)
+                             arguments_ty args, expr_ty returns)
 {
     /* Push arg annotation names and values.
        The expressions are evaluated separately from the rest of the source code.
@@ -1050,9 +1112,8 @@ codegen_annotations(compiler *c, location loc,
         RETURN_IF_ERROR_IN_SCOPE(
             c, codegen_annotations_in_scope(c, loc, args, returns, &annotations_len)
         );
-        RETURN_IF_ERROR(
+        ADDOP_I(c, loc, BUILD_MAP, annotations_len);
-            codegen_leave_annotations_scope(c, loc, annotations_len)
+        RETURN_IF_ERROR(codegen_leave_annotations_scope(c, loc));
-        );
         return MAKE_FUNCTION_ANNOTATE;
     }
     else {
@@ -1378,7 +1439,7 @@ codegen_function(compiler *c, stmt_ty s, int is_async)
         }
     }
 
-    int annotations_flag = codegen_annotations(c, loc, args, returns);
+    int annotations_flag = codegen_function_annotations(c, loc, args, returns);
     if (annotations_flag < 0) {
         if (is_generic) {
             _PyCompile_ExitScope(c);
@@ -1471,8 +1532,12 @@ codegen_class_body(compiler *c, stmt_ty s, int firstlineno)
         // that by default.
         ADDOP_N_IN_SCOPE(c, loc, STORE_DEREF, &_Py_ID(__classdict__), cellvars);
     }
+    if (SYMTABLE_ENTRY(c)->ste_has_conditional_annotations) {
+        ADDOP_I(c, loc, BUILD_SET, 0);
+        ADDOP_N_IN_SCOPE(c, loc, STORE_DEREF, &_Py_ID(__conditional_annotations__), cellvars);
+    }
     /* compile the body proper */
-    RETURN_IF_ERROR_IN_SCOPE(c, _PyCodegen_Body(c, loc, s->v.ClassDef.body, false));
+    RETURN_IF_ERROR_IN_SCOPE(c, codegen_body(c, loc, s->v.ClassDef.body, false));
     PyObject *static_attributes = _PyCompile_StaticAttributesAsTuple(c);
     if (static_attributes == NULL) {
         _PyCompile_ExitScope(c);
@@ -2895,6 +2960,14 @@ codegen_stmt_expr(compiler *c, location loc, expr_ty value)
     return SUCCESS;
 }
 
+#define CODEGEN_COND_BLOCK(FUNC, C, S) \
+    do { \
+        _PyCompile_EnterConditionalBlock((C)); \
+        int result = FUNC((C), (S)); \
+        _PyCompile_LeaveConditionalBlock((C)); \
+        return result; \
+    } while(0)
+
 static int
 codegen_visit_stmt(compiler *c, stmt_ty s)
 {
@@ -2929,13 +3002,17 @@ codegen_visit_stmt(compiler *c, stmt_ty s)
     case AnnAssign_kind:
         return codegen_annassign(c, s);
     case For_kind:
-        return codegen_for(c, s);
+        CODEGEN_COND_BLOCK(codegen_for, c, s);
+        break;
     case While_kind:
-        return codegen_while(c, s);
+        CODEGEN_COND_BLOCK(codegen_while, c, s);
+        break;
     case If_kind:
-        return codegen_if(c, s);
+        CODEGEN_COND_BLOCK(codegen_if, c, s);
+        break;
     case Match_kind:
-        return codegen_match(c, s);
+        CODEGEN_COND_BLOCK(codegen_match, c, s);
+        break;
     case Raise_kind:
     {
         Py_ssize_t n = 0;
@@ -2951,9 +3028,11 @@ codegen_visit_stmt(compiler *c, stmt_ty s)
         break;
     }
     case Try_kind:
-        return codegen_try(c, s);
+        CODEGEN_COND_BLOCK(codegen_try, c, s);
+        break;
     case TryStar_kind:
-        return codegen_try_star(c, s);
+        CODEGEN_COND_BLOCK(codegen_try_star, c, s);
+        break;
     case Assert_kind:
         return codegen_assert(c, s);
     case Import_kind:
@@ -2981,13 +3060,16 @@ codegen_visit_stmt(compiler *c, stmt_ty s)
         return codegen_continue(c, LOC(s));
     }
     case With_kind:
-        return codegen_with(c, s, 0);
+        CODEGEN_COND_BLOCK(codegen_with, c, s);
+        break;
     case AsyncFunctionDef_kind:
         return codegen_function(c, s, 1);
     case AsyncWith_kind:
-        return codegen_async_with(c, s, 0);
+        CODEGEN_COND_BLOCK(codegen_async_with, c, s);
+        break;
     case AsyncFor_kind:
-        return codegen_async_for(c, s);
+        CODEGEN_COND_BLOCK(codegen_async_for, c, s);
+        break;
     }
 
     return SUCCESS;
@@ -4725,7 +4807,7 @@ codegen_with_except_finish(compiler *c, jump_target_label cleanup) {
            raise
  */
 static int
-codegen_async_with(compiler *c, stmt_ty s, int pos)
+codegen_async_with_inner(compiler *c, stmt_ty s, int pos)
 {
     location loc = LOC(s);
     withitem_ty item = asdl_seq_GET(s->v.AsyncWith.items, pos);
@@ -4770,7 +4852,7 @@ codegen_async_with(compiler *c, stmt_ty s, int pos)
         VISIT_SEQ(c, stmt, s->v.AsyncWith.body);
     }
     else {
-        RETURN_IF_ERROR(codegen_async_with(c, s, pos));
+        RETURN_IF_ERROR(codegen_async_with_inner(c, s, pos));
     }
 
     _PyCompile_PopFBlock(c, COMPILE_FBLOCK_ASYNC_WITH, block);
@@ -4805,6 +4887,12 @@ codegen_async_with(compiler *c, stmt_ty s, int pos)
     return SUCCESS;
 }
 
+static int
+codegen_async_with(compiler *c, stmt_ty s)
+{
+    return codegen_async_with_inner(c, s, 0);
+}
+
 
 /*
    Implements the with statement from PEP 343.
@@ -4828,7 +4916,7 @@ codegen_async_with(compiler *c, stmt_ty s, int pos)
  */
 
 static int
-codegen_with(compiler *c, stmt_ty s, int pos)
+codegen_with_inner(compiler *c, stmt_ty s, int pos)
 {
     withitem_ty item = asdl_seq_GET(s->v.With.items, pos);
 
@@ -4869,7 +4957,7 @@ codegen_with(compiler *c, stmt_ty s, int pos)
         VISIT_SEQ(c, stmt, s->v.With.body);
     }
     else {
-        RETURN_IF_ERROR(codegen_with(c, s, pos));
+        RETURN_IF_ERROR(codegen_with_inner(c, s, pos));
     }
 
     ADDOP(c, NO_LOCATION, POP_BLOCK);
@@ -4896,6 +4984,12 @@ codegen_with(compiler *c, stmt_ty s, int pos)
     return SUCCESS;
 }
 
+static int
+codegen_with(compiler *c, stmt_ty s)
+{
+    return codegen_with_inner(c, s, 0);
+}
+
 static int
 codegen_visit_expr(compiler *c, expr_ty e)
 {
@@ -5224,7 +5318,18 @@ codegen_annassign(compiler *c, stmt_ty s)
                 ADDOP(c, loc, STORE_SUBSCR);
             }
             else {
-                RETURN_IF_ERROR(_PyCompile_AddDeferredAnnotaion(c, s));
+                PyObject *conditional_annotation_index = NULL;
+                RETURN_IF_ERROR(_PyCompile_AddDeferredAnnotation(
+                    c, s, &conditional_annotation_index));
+                if (conditional_annotation_index != NULL) {
+                    ADDOP_NAME(
+                        c, loc,
+                        SCOPE_TYPE(c) == COMPILE_SCOPE_CLASS ? LOAD_DEREF : LOAD_NAME,
+                        &_Py_ID(__conditional_annotations__), cellvars);
+                    ADDOP_LOAD_CONST_NEW(c, loc, conditional_annotation_index);
+                    ADDOP_I(c, loc, SET_ADD, 1);
+                    ADDOP(c, loc, POP_TOP);
+                }
             }
         }
         break;

Python/compile.c

@@ -60,11 +60,14 @@ struct compiler_unit {
     PyObject *u_private;            /* for private name mangling */
     PyObject *u_static_attributes;  /* for class: attributes accessed via self.X */
     PyObject *u_deferred_annotations; /* AnnAssign nodes deferred to the end of compilation */
+    PyObject *u_conditional_annotation_indices;  /* indices of annotations that are conditionally executed (or -1 for unconditional annotations) */
+    long u_next_conditional_annotation_index;  /* index of the next conditional annotation */
 
     instr_sequence *u_instr_sequence; /* codegen output */
 
     int u_nfblocks;
     int u_in_inlined_comp;
+    int u_in_conditional_block;
 
     _PyCompile_FBlockInfo u_fblock[CO_MAXBLOCKS];
 
@@ -187,6 +190,7 @@ compiler_unit_free(struct compiler_unit *u)
     Py_CLEAR(u->u_private);
     Py_CLEAR(u->u_static_attributes);
     Py_CLEAR(u->u_deferred_annotations);
+    Py_CLEAR(u->u_conditional_annotation_indices);
     PyMem_Free(u);
 }
 
@@ -620,6 +624,16 @@ _PyCompile_EnterScope(compiler *c, identifier name, int scope_type,
             return ERROR;
         }
     }
+    if (u->u_ste->ste_has_conditional_annotations) {
+        /* Cook up an implicit __conditional__annotations__ cell */
+        Py_ssize_t res;
+        assert(u->u_scope_type == COMPILE_SCOPE_CLASS || u->u_scope_type == COMPILE_SCOPE_MODULE);
+        res = _PyCompile_DictAddObj(u->u_metadata.u_cellvars, &_Py_ID(__conditional_annotations__));
+        if (res < 0) {
+            compiler_unit_free(u);
+            return ERROR;
+        }
+    }
 
     u->u_metadata.u_freevars = dictbytype(u->u_ste->ste_symbols, FREE, DEF_FREE_CLASS,
                                PyDict_GET_SIZE(u->u_metadata.u_cellvars));
@@ -649,6 +663,8 @@ _PyCompile_EnterScope(compiler *c, identifier name, int scope_type,
     }
 
     u->u_deferred_annotations = NULL;
+    u->u_conditional_annotation_indices = NULL;
+    u->u_next_conditional_annotation_index = 0;
     if (scope_type == COMPILE_SCOPE_CLASS) {
         u->u_static_attributes = PySet_New(0);
         if (!u->u_static_attributes) {
@@ -768,10 +784,13 @@ _PyCompile_TopFBlock(compiler *c)
     return &c->u->u_fblock[c->u->u_nfblocks - 1];
 }
 
-PyObject *
+void
-_PyCompile_DeferredAnnotations(compiler *c)
+_PyCompile_DeferredAnnotations(compiler *c,
+                               PyObject **deferred_annotations,
+                               PyObject **conditional_annotation_indices)
 {
-    return c->u->u_deferred_annotations;
+    *deferred_annotations = Py_XNewRef(c->u->u_deferred_annotations);
+    *conditional_annotation_indices = Py_XNewRef(c->u->u_conditional_annotation_indices);
 }
 
 static location
@@ -797,13 +816,13 @@ compiler_codegen(compiler *c, mod_ty mod)
     switch (mod->kind) {
     case Module_kind: {
         asdl_stmt_seq *stmts = mod->v.Module.body;
-        RETURN_IF_ERROR(_PyCodegen_Body(c, start_location(stmts), stmts, false));
+        RETURN_IF_ERROR(_PyCodegen_Module(c, start_location(stmts), stmts, false));
         break;
     }
     case Interactive_kind: {
         c->c_interactive = 1;
         asdl_stmt_seq *stmts = mod->v.Interactive.body;
-        RETURN_IF_ERROR(_PyCodegen_Body(c, start_location(stmts), stmts, true));
+        RETURN_IF_ERROR(_PyCodegen_Module(c, start_location(stmts), stmts, true));
         break;
     }
     case Expression_kind: {
@@ -838,7 +857,8 @@ _PyCompile_GetRefType(compiler *c, PyObject *name)
 {
     if (c->u->u_scope_type == COMPILE_SCOPE_CLASS &&
         (_PyUnicode_EqualToASCIIString(name, "__class__") ||
-         _PyUnicode_EqualToASCIIString(name, "__classdict__"))) {
+         _PyUnicode_EqualToASCIIString(name, "__classdict__") ||
+         _PyUnicode_EqualToASCIIString(name, "__conditional_annotations__"))) {
         return CELL;
     }
     PySTEntryObject *ste = c->u->u_ste;
@@ -1095,8 +1115,22 @@ _PyCompile_RevertInlinedComprehensionScopes(compiler *c, location loc,
     return SUCCESS;
 }
 
+void
+_PyCompile_EnterConditionalBlock(struct _PyCompiler *c)
+{
+    c->u->u_in_conditional_block++;
+}
+
+void
+_PyCompile_LeaveConditionalBlock(struct _PyCompiler *c)
+{
+    assert(c->u->u_in_conditional_block > 0);
+    c->u->u_in_conditional_block--;
+}
+
 int
-_PyCompile_AddDeferredAnnotaion(compiler *c, stmt_ty s)
+_PyCompile_AddDeferredAnnotation(compiler *c, stmt_ty s,
+                                 PyObject **conditional_annotation_index)
 {
     if (c->u->u_deferred_annotations == NULL) {
         c->u->u_deferred_annotations = PyList_New(0);
@@ -1104,6 +1138,12 @@ _PyCompile_AddDeferredAnnotaion(compiler *c, stmt_ty s)
             return ERROR;
         }
     }
+    if (c->u->u_conditional_annotation_indices == NULL) {
+        c->u->u_conditional_annotation_indices = PyList_New(0);
+        if (c->u->u_conditional_annotation_indices == NULL) {
+            return ERROR;
+        }
+    }
     PyObject *ptr = PyLong_FromVoidPtr((void *)s);
     if (ptr == NULL) {
         return ERROR;
@@ -1113,6 +1153,24 @@ _PyCompile_AddDeferredAnnotaion(compiler *c, stmt_ty s)
         return ERROR;
     }
     Py_DECREF(ptr);
+    PyObject *index;
+    if (c->u->u_in_conditional_block) {
+        index = PyLong_FromLong(c->u->u_next_conditional_annotation_index);
+        if (index == NULL) {
+            return ERROR;
+        }
+        *conditional_annotation_index = Py_NewRef(index);
+        c->u->u_next_conditional_annotation_index++;
+    }
+    else {
+        index = PyLong_FromLong(-1);
+        if (index == NULL) {
+            return ERROR;
+        }
+    }
+    int rc = PyList_Append(c->u->u_conditional_annotation_indices, index);
+    Py_DECREF(index);
+    RETURN_IF_ERROR(rc);
     return SUCCESS;
 }
 

Python/symtable.c

@@ -139,6 +139,8 @@ ste_new(struct symtable *st, identifier name, _Py_block_ty block,
     ste->ste_can_see_class_scope = 0;
     ste->ste_comp_iter_expr = 0;
     ste->ste_needs_classdict = 0;
+    ste->ste_has_conditional_annotations = 0;
+    ste->ste_in_conditional_block = 0;
     ste->ste_annotation_block = NULL;
 
     ste->ste_has_docstring = 0;
@@ -945,6 +947,12 @@ drop_class_free(PySTEntryObject *ste, PyObject *free)
         return 0;
     if (res)
         ste->ste_needs_classdict = 1;
+    res = PySet_Discard(free, &_Py_ID(__conditional_annotations__));
+    if (res < 0)
+        return 0;
+    if (res) {
+        ste->ste_has_conditional_annotations = 1;
+    }
     return 1;
 }
 
@@ -1198,6 +1206,8 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
             goto error;
         if (PySet_Add(newbound, &_Py_ID(__classdict__)) < 0)
             goto error;
+        if (PySet_Add(newbound, &_Py_ID(__conditional_annotations__)) < 0)
+            goto error;
     }
 
     /* Recursively call analyze_child_block() on each child block.
@@ -1729,6 +1739,13 @@ symtable_enter_type_param_block(struct symtable *st, identifier name,
         } \
     } while(0)
 
+#define ENTER_CONDITIONAL_BLOCK(ST) \
+    int in_conditional_block = (ST)->st_cur->ste_in_conditional_block; \
+    (ST)->st_cur->ste_in_conditional_block = 1;
+
+#define LEAVE_CONDITIONAL_BLOCK(ST) \
+    (ST)->st_cur->ste_in_conditional_block = in_conditional_block;
+
 #define ENTER_RECURSIVE() \
 if (Py_EnterRecursiveCall(" during compilation")) { \
     return 0; \
@@ -2012,30 +2029,42 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
         VISIT(st, expr, s->v.AugAssign.value);
         break;
     }
-    case For_kind:
+    case For_kind: {
         VISIT(st, expr, s->v.For.target);
         VISIT(st, expr, s->v.For.iter);
+        ENTER_CONDITIONAL_BLOCK(st);
         VISIT_SEQ(st, stmt, s->v.For.body);
         if (s->v.For.orelse)
             VISIT_SEQ(st, stmt, s->v.For.orelse);
+        LEAVE_CONDITIONAL_BLOCK(st);
         break;
-    case While_kind:
+    }
+    case While_kind: {
         VISIT(st, expr, s->v.While.test);
+        ENTER_CONDITIONAL_BLOCK(st);
         VISIT_SEQ(st, stmt, s->v.While.body);
         if (s->v.While.orelse)
             VISIT_SEQ(st, stmt, s->v.While.orelse);
+        LEAVE_CONDITIONAL_BLOCK(st);
         break;
-    case If_kind:
+    }
+    case If_kind: {
         /* XXX if 0: and lookup_yield() hacks */
         VISIT(st, expr, s->v.If.test);
+        ENTER_CONDITIONAL_BLOCK(st);
         VISIT_SEQ(st, stmt, s->v.If.body);
         if (s->v.If.orelse)
             VISIT_SEQ(st, stmt, s->v.If.orelse);
+        LEAVE_CONDITIONAL_BLOCK(st);
         break;
-    case Match_kind:
+    }
+    case Match_kind: {
         VISIT(st, expr, s->v.Match.subject);
+        ENTER_CONDITIONAL_BLOCK(st);
         VISIT_SEQ(st, match_case, s->v.Match.cases);
+        LEAVE_CONDITIONAL_BLOCK(st);
         break;
+    }
     case Raise_kind:
         if (s->v.Raise.exc) {
             VISIT(st, expr, s->v.Raise.exc);
@@ -2044,18 +2073,24 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
             }
         }
         break;
-    case Try_kind:
+    case Try_kind: {
+        ENTER_CONDITIONAL_BLOCK(st);
         VISIT_SEQ(st, stmt, s->v.Try.body);
         VISIT_SEQ(st, excepthandler, s->v.Try.handlers);
         VISIT_SEQ(st, stmt, s->v.Try.orelse);
         VISIT_SEQ(st, stmt, s->v.Try.finalbody);
+        LEAVE_CONDITIONAL_BLOCK(st);
         break;
-    case TryStar_kind:
+    }
+    case TryStar_kind: {
+        ENTER_CONDITIONAL_BLOCK(st);
         VISIT_SEQ(st, stmt, s->v.TryStar.body);
         VISIT_SEQ(st, excepthandler, s->v.TryStar.handlers);
         VISIT_SEQ(st, stmt, s->v.TryStar.orelse);
         VISIT_SEQ(st, stmt, s->v.TryStar.finalbody);
+        LEAVE_CONDITIONAL_BLOCK(st);
         break;
+    }
     case Assert_kind:
         VISIT(st, expr, s->v.Assert.test);
         if (s->v.Assert.msg)
@@ -2142,10 +2177,13 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
     case Continue_kind:
         /* nothing to do here */
         break;
-    case With_kind:
+    case With_kind: {
+        ENTER_CONDITIONAL_BLOCK(st);
         VISIT_SEQ(st, withitem, s->v.With.items);
         VISIT_SEQ(st, stmt, s->v.With.body);
+        LEAVE_CONDITIONAL_BLOCK(st);
         break;
+    }
     case AsyncFunctionDef_kind: {
         if (!symtable_add_def(st, s->v.AsyncFunctionDef.name, DEF_LOCAL, LOCATION(s)))
             return 0;
@@ -2201,26 +2239,32 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
         }
         break;
     }
-    case AsyncWith_kind:
+    case AsyncWith_kind: {
         maybe_set_ste_coroutine_for_module(st, s);
         if (!symtable_raise_if_not_coroutine(st, ASYNC_WITH_OUTSIDE_ASYNC_FUNC, LOCATION(s))) {
             return 0;
         }
+        ENTER_CONDITIONAL_BLOCK(st);
         VISIT_SEQ(st, withitem, s->v.AsyncWith.items);
         VISIT_SEQ(st, stmt, s->v.AsyncWith.body);
+        LEAVE_CONDITIONAL_BLOCK(st);
         break;
-    case AsyncFor_kind:
+    }
+    case AsyncFor_kind: {
         maybe_set_ste_coroutine_for_module(st, s);
         if (!symtable_raise_if_not_coroutine(st, ASYNC_FOR_OUTSIDE_ASYNC_FUNC, LOCATION(s))) {
             return 0;
         }
         VISIT(st, expr, s->v.AsyncFor.target);
         VISIT(st, expr, s->v.AsyncFor.iter);
+        ENTER_CONDITIONAL_BLOCK(st);
         VISIT_SEQ(st, stmt, s->v.AsyncFor.body);
         if (s->v.AsyncFor.orelse)
             VISIT_SEQ(st, stmt, s->v.AsyncFor.orelse);
+        LEAVE_CONDITIONAL_BLOCK(st);
         break;
     }
+    }
     LEAVE_RECURSIVE();
     return 1;
 }
@@ -2689,6 +2733,15 @@ symtable_visit_params(struct symtable *st, asdl_arg_seq *args)
 static int
 symtable_visit_annotation(struct symtable *st, expr_ty annotation, void *key)
 {
+    if ((st->st_cur->ste_type == ClassBlock || st->st_cur->ste_type == ModuleBlock)
+            && st->st_cur->ste_in_conditional_block
+            && !st->st_cur->ste_has_conditional_annotations)
+    {
+        st->st_cur->ste_has_conditional_annotations = 1;
+        if (!symtable_add_def(st, &_Py_ID(__conditional_annotations__), USE, LOCATION(annotation))) {
+            return 0;
+        }
+    }
     struct _symtable_entry *parent_ste = st->st_cur;
     if (parent_ste->ste_annotation_block == NULL) {
         _Py_block_ty current_type = parent_ste->ste_type;