Treat all module-level annotations as conditional

Author: JelleZijlstra

Include/internal/pycore_compile.h

@@ -133,6 +133,8 @@ int _PyCompile_EnterScope(struct _PyCompiler *c, identifier name, int scope_type
 void _PyCompile_ExitScope(struct _PyCompiler *c);
 Py_ssize_t _PyCompile_AddConst(struct _PyCompiler *c, PyObject *o);
 _PyInstructionSequence *_PyCompile_InstrSequence(struct _PyCompiler *c);
+void _PyCompile_SetInstrSequence(struct _PyCompiler *c,
+                                 _PyInstructionSequence *instr_sequence);
 int _PyCompile_FutureFeatures(struct _PyCompiler *c);
 void _PyCompile_DeferredAnnotations(
     struct _PyCompiler *c, PyObject **deferred_annotations,

Include/internal/pycore_instruction_sequence.h

@@ -66,6 +66,7 @@ _PyJumpTargetLabel _PyInstructionSequence_NewLabel(_PyInstructionSequence *seq);
 int _PyInstructionSequence_ApplyLabelMap(_PyInstructionSequence *seq);
 int _PyInstructionSequence_InsertInstruction(_PyInstructionSequence *seq, int pos,
                                              int opcode, int oparg, _Py_SourceLocation loc);
+int _PyInstructionSequence_PrependSequence(_PyInstructionSequence *seq, _PyInstructionSequence *nested);
 int _PyInstructionSequence_AddNested(_PyInstructionSequence *seq, _PyInstructionSequence *nested);
 void PyInstructionSequence_Fini(_PyInstructionSequence *seq);
 

Lib/test/test_type_annotations.py

@@ -115,11 +115,7 @@ def test_partially_executed_module(self):
             partialexe.a.__annotations__,
             {"v1": int, "v2": int},
         )
-        self.assertIsInstance(partialexe.b.exc, RuntimeError)
-        self.assertEqual(
-            str(partialexe.b.exc),
-            "cannot access __annotations__ while module is initializing",
-        )
+        self.assertEqual(partialexe.b.annos, {"v1": int})
 
     @cpython_only
     def test_no_cell(self):

Lib/test/typinganndata/partialexecution/b.py

@@ -1,8 +1,3 @@
 from . import a
 
-try:
-    a.__annotations__
-except Exception as e:
-    exc = e
-else:
-    exc = None
+annos = a.__annotations__

Objects/moduleobject.c

@@ -1251,6 +1251,7 @@ module_get_annotations(PyObject *self, void *Py_UNUSED(ignored))
             Py_DECREF(dict);
             return NULL;
         }
+        bool is_initializing = false;
         if (spec != NULL) {
             int rc = _PyModuleSpec_IsInitializing(spec);
             if (rc < 0) {
@@ -1260,11 +1261,7 @@ module_get_annotations(PyObject *self, void *Py_UNUSED(ignored))
             }
             Py_DECREF(spec);
             if (rc) {
-                PyErr_SetString(PyExc_RuntimeError,
-                                "cannot access __annotations__ "
-                                "while module is initializing");
-                Py_DECREF(dict);
-                return NULL;
+                is_initializing = true;
             }
         }
 
@@ -1295,7 +1292,8 @@ module_get_annotations(PyObject *self, void *Py_UNUSED(ignored))
             annotations = PyDict_New();
         }
         Py_XDECREF(annotate);
-        if (annotations) {
+        // Do not cache annotations if the module is still initializing
+        if (annotations && !is_initializing) {
             int result = PyDict_SetItem(
                     dict, &_Py_ID(__annotations__), annotations);
             if (result) {

Python/codegen.c

@@ -792,6 +792,17 @@ codegen_process_deferred_annotations(compiler *c, location loc)
         return SUCCESS;
     }
 
+    instr_sequence *old_instr_seq = INSTR_SEQUENCE(c);
+    instr_sequence *nested_instr_seq = NULL;
+    int scope_type = SCOPE_TYPE(c);
+    if (scope_type == COMPILE_SCOPE_MODULE) {
+        nested_instr_seq = (instr_sequence *)_PyInstructionSequence_New();
+        if (nested_instr_seq == NULL) {
+            goto error;
+        }
+        _PyCompile_SetInstrSequence(c, nested_instr_seq);
+    }
+
     // It's possible that ste_annotations_block is set but
     // u_deferred_annotations is not, because the former is still
     // set if there are only non-simple annotations (i.e., annotations
@@ -800,7 +811,6 @@ 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) {
         goto error;
@@ -817,8 +827,19 @@ codegen_process_deferred_annotations(compiler *c, location loc)
     RETURN_IF_ERROR(codegen_leave_annotations_scope(c, loc));
     RETURN_IF_ERROR(codegen_nameop(c, loc, &_Py_ID(__annotate__), Store));
 
+    if (nested_instr_seq != NULL) {
+        RETURN_IF_ERROR(
+            _PyInstructionSequence_PrependSequence(old_instr_seq, nested_instr_seq));
+        _PyCompile_SetInstrSequence(c, old_instr_seq);
+        PyInstructionSequence_Fini(nested_instr_seq);
+    }
+
     return SUCCESS;
 error:
+    if (nested_instr_seq != NULL) {
+        PyInstructionSequence_Fini(nested_instr_seq);
+        _PyCompile_SetInstrSequence(c, old_instr_seq);
+    }
     Py_XDECREF(deferred_anno);
     Py_XDECREF(conditional_annotation_indices);
     return ERROR;

Python/compile.c

@@ -1154,7 +1154,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;
@@ -1231,6 +1231,12 @@ _PyCompile_InstrSequence(compiler *c)
     return c->u->u_instr_sequence;
 }
 
+void
+_PyCompile_SetInstrSequence(compiler *c, instr_sequence *seq)
+{
+    c->u->u_instr_sequence = seq;
+}
+
 int
 _PyCompile_FutureFeatures(compiler *c)
 {

Python/instruction_sequence.c

@@ -34,24 +34,29 @@ typedef _Py_SourceLocation location;
     }
 
 static int
-instr_sequence_next_inst(instr_sequence *seq) {
+instr_sequence_grow(instr_sequence *seq, int size) {
     assert(seq->s_instrs != NULL || seq->s_used == 0);
 
-
     _Py_c_array_t array = {
         .array = (void*)seq->s_instrs,
         .allocated_entries = seq->s_allocated,
         .item_size = sizeof(instruction),
         .initial_num_entries = INITIAL_INSTR_SEQUENCE_SIZE,
     };
 
-    RETURN_IF_ERROR(_Py_CArray_EnsureCapacity(&array, seq->s_used + 1));
+    RETURN_IF_ERROR(_Py_CArray_EnsureCapacity(&array, seq->s_used + size));
     seq->s_instrs = array.array;
     seq->s_allocated = array.allocated_entries;
 
     assert(seq->s_allocated >= 0);
     assert(seq->s_used < seq->s_allocated);
-    return seq->s_used++;
+    seq->s_used += size;
+    return seq->s_used - 1;
+}
+
+static int
+instr_sequence_next_inst(instr_sequence *seq) {
+    return instr_sequence_grow(seq, 1);
 }
 
 _PyJumpTargetLabel
@@ -154,6 +159,31 @@ _PyInstructionSequence_InsertInstruction(instr_sequence *seq, int pos,
     return SUCCESS;
 }
 
+int
+_PyInstructionSequence_PrependSequence(instr_sequence *seq,
+                                       instr_sequence *nested)
+{
+    if (nested->s_used == 0) {
+        return SUCCESS;
+    }
+    // Merging labelmaps is not supported
+    assert(nested->s_labelmap_size == 0 && nested->s_nested == NULL);
+
+    int last_idx = instr_sequence_grow(seq, nested->s_used);
+
+    RETURN_IF_ERROR(last_idx);
+    for (int i = last_idx - nested->s_used; i >= 0; i--) {
+        seq->s_instrs[i + nested->s_used] = seq->s_instrs[i];
+    }
+    for (int i=0; i < nested->s_used; i++) {
+        seq->s_instrs[i] = nested->s_instrs[i];
+    }
+    for(int lbl=0; lbl < seq->s_labelmap_size; lbl++) {
+        seq->s_labelmap[lbl] += nested->s_used;
+    }
+    return SUCCESS;
+}
+
 int
 _PyInstructionSequence_AddNested(instr_sequence *seq, instr_sequence *nested)
 {

Python/symtable.c

@@ -2749,8 +2749,10 @@ symtable_visit_annotation(struct symtable *st, expr_ty annotation, void *key)
     // Annotations in local scopes are not executed and should not affect the symtable
     bool is_unevaluated = st->st_cur->ste_type == FunctionBlock;
 
-    if ((st->st_cur->ste_type == ClassBlock || st->st_cur->ste_type == ModuleBlock)
-            && st->st_cur->ste_in_conditional_block
+    // Module-level annotations are always considered conditional because the module
+    // may be partially executed.
+    if ((((st->st_cur->ste_type == ClassBlock && st->st_cur->ste_in_conditional_block)
+            || st->st_cur->ste_type == ModuleBlock))
             && !st->st_cur->ste_has_conditional_annotations)
     {
         st->st_cur->ste_has_conditional_annotations = 1;