gh-130907: Treat all module-level annotations as conditional (#131550)

1 files changed, 33 insertions, 1 deletions

diff --git a/Python/compile.c b/Python/compile.c
index 430898b..15ef721 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -64,6 +64,7 @@ struct compiler_unit {
     long u_next_conditional_annotation_index;  /* index of the next conditional annotation */
 
     instr_sequence *u_instr_sequence; /* codegen output */
+    instr_sequence *u_stashed_instr_sequence; /* temporarily stashed parent instruction sequence */
 
     int u_nfblocks;
     int u_in_inlined_comp;
@@ -178,6 +179,7 @@ static void
 compiler_unit_free(struct compiler_unit *u)
 {
     Py_CLEAR(u->u_instr_sequence);
+    Py_CLEAR(u->u_stashed_instr_sequence);
     Py_CLEAR(u->u_ste);
     Py_CLEAR(u->u_metadata.u_name);
     Py_CLEAR(u->u_metadata.u_qualname);
@@ -681,6 +683,7 @@ _PyCompile_EnterScope(compiler *c, identifier name, int scope_type,
         compiler_unit_free(u);
         return ERROR;
     }
+    u->u_stashed_instr_sequence = NULL;
 
     /* Push the old compiler_unit on the stack. */
     if (c->u) {
@@ -1154,7 +1157,7 @@ _PyCompile_AddDeferredAnnotation(compiler *c, stmt_ty s,
     }
     Py_DECREF(ptr);
     PyObject *index;
-    if (c->u->u_in_conditional_block) {
+    if (c->u->u_scope_type == COMPILE_SCOPE_MODULE || c->u->u_in_conditional_block) {
         index = PyLong_FromLong(c->u->u_next_conditional_annotation_index);
         if (index == NULL) {
             return ERROR;
@@ -1232,6 +1235,35 @@ _PyCompile_InstrSequence(compiler *c)
 }
 
 int
+_PyCompile_StartAnnotationSetup(struct _PyCompiler *c)
+{
+    instr_sequence *new_seq = (instr_sequence *)_PyInstructionSequence_New();
+    if (new_seq == NULL) {
+        return ERROR;
+    }
+    assert(c->u->u_stashed_instr_sequence == NULL);
+    c->u->u_stashed_instr_sequence = c->u->u_instr_sequence;
+    c->u->u_instr_sequence = new_seq;
+    return SUCCESS;
+}
+
+int
+_PyCompile_EndAnnotationSetup(struct _PyCompiler *c)
+{
+    assert(c->u->u_stashed_instr_sequence != NULL);
+    instr_sequence *parent_seq = c->u->u_stashed_instr_sequence;
+    instr_sequence *anno_seq = c->u->u_instr_sequence;
+    c->u->u_stashed_instr_sequence = NULL;
+    c->u->u_instr_sequence = parent_seq;
+    if (_PyInstructionSequence_SetAnnotationsCode(parent_seq, anno_seq) == ERROR) {
+        Py_DECREF(anno_seq);
+        return ERROR;
+    }
+    return SUCCESS;
+}
+
+
+int
 _PyCompile_FutureFeatures(compiler *c)
 {
     return c->c_future.ff_features;