gh-130881: Handle conditionally defined annotations

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_strings.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_symtable.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 */

Lib/test/test_type_annotations.py

@@ -457,3 +457,202 @@ 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 (
+                # 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.