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

This commit is contained in:
Jelle Zijlstra 2025-03-25 20:48:19 -07:00 committed by GitHub
parent 7bb41aef4b
commit 898e6b395e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 506 additions and 74 deletions

View file

@ -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 *
_PyCompile_DeferredAnnotations(compiler *c)
void
_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;
}