Skip to content

Commit e077b20

Browse files
iritkatrieliritkatriel
authored
gh-121404: pass metadata to compiler_enter_scope (optionally) so that codegen functions don't need to set it on the code unit (#123078)
1 parent b9e10d1 commit e077b20

File tree

1 file changed

+82
-58
lines changed

1 file changed

+82
-58
lines changed

‎Python/compile.c‎

Lines changed: 82 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -1057,7 +1057,8 @@ codegen_addop_j(instr_sequence *seq, location loc,
10571057

10581058
static int
10591059
compiler_enter_scope(struct compiler *c, identifier name, int scope_type,
1060-
void *key, int lineno, PyObject *private)
1060+
void *key, int lineno, PyObject *private,
1061+
_PyCompile_CodeUnitMetadata *umd)
10611062
{
10621063
location loc = LOCATION(lineno, lineno, 0, 0);
10631064

@@ -1069,9 +1070,14 @@ compiler_enter_scope(struct compiler *c, identifier name, int scope_type,
10691070
return ERROR;
10701071
}
10711072
u->u_scope_type = scope_type;
1072-
u->u_metadata.u_argcount = 0;
1073-
u->u_metadata.u_posonlyargcount = 0;
1074-
u->u_metadata.u_kwonlyargcount = 0;
1073+
if (umd != NULL) {
1074+
u->u_metadata = *umd;
1075+
}
1076+
else {
1077+
u->u_metadata.u_argcount = 0;
1078+
u->u_metadata.u_posonlyargcount = 0;
1079+
u->u_metadata.u_kwonlyargcount = 0;
1080+
}
10751081
u->u_ste = _PySymtable_Lookup(c->c_st, key);
10761082
if (!u->u_ste) {
10771083
compiler_unit_free(u);
@@ -1441,14 +1447,19 @@ compiler_unwind_fblock_stack(struct compiler *c, location *ploc,
14411447
}
14421448

14431449
static int
1444-
compiler_setup_annotations_scope(struct compiler *c, location loc,
1445-
void *key, PyObject *name)
1450+
codegen_setup_annotations_scope(struct compiler *c, location loc,
1451+
void *key, PyObject *name)
14461452
{
1447-
if (compiler_enter_scope(c, name, COMPILER_SCOPE_ANNOTATIONS,
1448-
key, loc.lineno, NULL) == -1) {
1449-
return ERROR;
1450-
}
1451-
c->u->u_metadata.u_posonlyargcount = 1;
1453+
_PyCompile_CodeUnitMetadata umd = {
1454+
.u_posonlyargcount = 1,
1455+
};
1456+
RETURN_IF_ERROR(
1457+
compiler_enter_scope(c, name, COMPILER_SCOPE_ANNOTATIONS,
1458+
key, loc.lineno, NULL, &umd));
1459+
assert(c->u->u_metadata.u_posonlyargcount == 1);
1460+
assert(c->u->u_metadata.u_argcount == 0);
1461+
assert(c->u->u_metadata.u_kwonlyargcount == 0);
1462+
14521463
// if .format != 1: raise NotImplementedError
14531464
_Py_DECLARE_STR(format, ".format");
14541465
ADDOP_I(c, loc, LOAD_FAST, 0);
@@ -1463,8 +1474,8 @@ compiler_setup_annotations_scope(struct compiler *c, location loc,
14631474
}
14641475

14651476
static int
1466-
compiler_leave_annotations_scope(struct compiler *c, location loc,
1467-
Py_ssize_t annotations_len)
1477+
codegen_leave_annotations_scope(struct compiler *c, location loc,
1478+
Py_ssize_t annotations_len)
14681479
{
14691480
ADDOP_I(c, loc, BUILD_MAP, annotations_len);
14701481
ADDOP_IN_SCOPE(c, loc, RETURN_VALUE);
@@ -1534,8 +1545,8 @@ compiler_body(struct compiler *c, location loc, asdl_stmt_seq *stmts)
15341545
assert(ste->ste_annotation_block != NULL);
15351546
PyObject *deferred_anno = Py_NewRef(c->u->u_deferred_annotations);
15361547
void *key = (void *)((uintptr_t)ste->ste_id + 1);
1537-
if (compiler_setup_annotations_scope(c, loc, key,
1538-
ste->ste_annotation_block->ste_name) == -1) {
1548+
if (codegen_setup_annotations_scope(c, loc, key,
1549+
ste->ste_annotation_block->ste_name) == -1) {
15391550
Py_DECREF(deferred_anno);
15401551
return ERROR;
15411552
}
@@ -1555,7 +1566,7 @@ compiler_body(struct compiler *c, location loc, asdl_stmt_seq *stmts)
15551566
Py_DECREF(deferred_anno);
15561567

15571568
RETURN_IF_ERROR(
1558-
compiler_leave_annotations_scope(c, loc, annotations_len)
1569+
codegen_leave_annotations_scope(c, loc, annotations_len)
15591570
);
15601571
RETURN_IF_ERROR(
15611572
compiler_nameop(c, loc, &_Py_ID(__annotate__), Store)
@@ -1618,7 +1629,7 @@ compiler_enter_anonymous_scope(struct compiler* c, mod_ty mod)
16181629
_Py_DECLARE_STR(anon_module, "<module>");
16191630
RETURN_IF_ERROR(
16201631
compiler_enter_scope(c, &_Py_STR(anon_module), COMPILER_SCOPE_MODULE,
1621-
mod, 1, NULL));
1632+
mod, 1, NULL, NULL));
16221633
return SUCCESS;
16231634
}
16241635

@@ -1828,7 +1839,7 @@ compiler_visit_annexpr(struct compiler *c, expr_ty annotation)
18281839
}
18291840

18301841
static int
1831-
compiler_argannotation(struct compiler *c, identifier id,
1842+
codegen_argannotation(struct compiler *c, identifier id,
18321843
expr_ty annotation, Py_ssize_t *annotations_len, location loc)
18331844
{
18341845
if (!annotation) {
@@ -1862,14 +1873,14 @@ compiler_argannotation(struct compiler *c, identifier id,
18621873
}
18631874

18641875
static int
1865-
compiler_argannotations(struct compiler *c, asdl_arg_seq* args,
1866-
Py_ssize_t *annotations_len, location loc)
1876+
codegen_argannotations(struct compiler *c, asdl_arg_seq* args,
1877+
Py_ssize_t *annotations_len, location loc)
18671878
{
18681879
int i;
18691880
for (i = 0; i < asdl_seq_LEN(args); i++) {
18701881
arg_ty arg = (arg_ty)asdl_seq_GET(args, i);
18711882
RETURN_IF_ERROR(
1872-
compiler_argannotation(
1883+
codegen_argannotation(
18731884
c,
18741885
arg->arg,
18751886
arg->annotation,
@@ -1880,40 +1891,40 @@ compiler_argannotations(struct compiler *c, asdl_arg_seq* args,
18801891
}
18811892

18821893
static int
1883-
compiler_annotations_in_scope(struct compiler *c, location loc,
1884-
arguments_ty args, expr_ty returns,
1885-
Py_ssize_t *annotations_len)
1894+
codegen_annotations_in_scope(struct compiler *c, location loc,
1895+
arguments_ty args, expr_ty returns,
1896+
Py_ssize_t *annotations_len)
18861897
{
18871898
RETURN_IF_ERROR(
1888-
compiler_argannotations(c, args->args, annotations_len, loc));
1899+
codegen_argannotations(c, args->args, annotations_len, loc));
18891900

18901901
RETURN_IF_ERROR(
1891-
compiler_argannotations(c, args->posonlyargs, annotations_len, loc));
1902+
codegen_argannotations(c, args->posonlyargs, annotations_len, loc));
18921903

18931904
if (args->vararg && args->vararg->annotation) {
18941905
RETURN_IF_ERROR(
1895-
compiler_argannotation(c, args->vararg->arg,
1896-
args->vararg->annotation, annotations_len, loc));
1906+
codegen_argannotation(c, args->vararg->arg,
1907+
args->vararg->annotation, annotations_len, loc));
18971908
}
18981909

18991910
RETURN_IF_ERROR(
1900-
compiler_argannotations(c, args->kwonlyargs, annotations_len, loc));
1911+
codegen_argannotations(c, args->kwonlyargs, annotations_len, loc));
19011912

19021913
if (args->kwarg && args->kwarg->annotation) {
19031914
RETURN_IF_ERROR(
1904-
compiler_argannotation(c, args->kwarg->arg,
1905-
args->kwarg->annotation, annotations_len, loc));
1915+
codegen_argannotation(c, args->kwarg->arg,
1916+
args->kwarg->annotation, annotations_len, loc));
19061917
}
19071918

19081919
RETURN_IF_ERROR(
1909-
compiler_argannotation(c, &_Py_ID(return), returns, annotations_len, loc));
1920+
codegen_argannotation(c, &_Py_ID(return), returns, annotations_len, loc));
19101921

19111922
return 0;
19121923
}
19131924

19141925
static int
1915-
compiler_annotations(struct compiler *c, location loc,
1916-
arguments_ty args, expr_ty returns)
1926+
codegen_annotations(struct compiler *c, location loc,
1927+
arguments_ty args, expr_ty returns)
19171928
{
19181929
/* Push arg annotation names and values.
19191930
The expressions are evaluated separately from the rest of the source code.
@@ -1930,15 +1941,15 @@ compiler_annotations(struct compiler *c, location loc,
19301941
bool annotations_used = ste->ste_annotations_used;
19311942

19321943
if (annotations_used) {
1933-
if (compiler_setup_annotations_scope(c, loc, (void *)args,
1934-
ste->ste_name) < 0) {
1944+
if (codegen_setup_annotations_scope(c, loc, (void *)args,
1945+
ste->ste_name) < 0) {
19351946
Py_DECREF(ste);
19361947
return ERROR;
19371948
}
19381949
}
19391950
Py_DECREF(ste);
19401951

1941-
if (compiler_annotations_in_scope(c, loc, args, returns, &annotations_len) < 0) {
1952+
if (codegen_annotations_in_scope(c, loc, args, returns, &annotations_len) < 0) {
19421953
if (annotations_used) {
19431954
compiler_exit_scope(c);
19441955
}
@@ -1947,7 +1958,7 @@ compiler_annotations(struct compiler *c, location loc,
19471958

19481959
if (annotations_used) {
19491960
RETURN_IF_ERROR(
1950-
compiler_leave_annotations_scope(c, loc, annotations_len)
1961+
codegen_leave_annotations_scope(c, loc, annotations_len)
19511962
);
19521963
return MAKE_FUNCTION_ANNOTATE;
19531964
}
@@ -2011,7 +2022,7 @@ compiler_type_param_bound_or_default(struct compiler *c, expr_ty e,
20112022
{
20122023
PyObject *defaults = PyTuple_Pack(1, _PyLong_GetOne());
20132024
ADDOP_LOAD_CONST_NEW(c, LOC(e), defaults);
2014-
if (compiler_setup_annotations_scope(c, LOC(e), key, name) == -1) {
2025+
if (codegen_setup_annotations_scope(c, LOC(e), key, name) == -1) {
20152026
return ERROR;
20162027
}
20172028
if (allow_starred && e->kind == Starred_kind) {
@@ -2155,8 +2166,13 @@ compiler_function_body(struct compiler *c, stmt_ty s, int is_async, Py_ssize_t f
21552166
scope_type = COMPILER_SCOPE_FUNCTION;
21562167
}
21572168

2169+
_PyCompile_CodeUnitMetadata umd = {
2170+
.u_argcount = asdl_seq_LEN(args->args),
2171+
.u_posonlyargcount = asdl_seq_LEN(args->posonlyargs),
2172+
.u_kwonlyargcount = asdl_seq_LEN(args->kwonlyargs),
2173+
};
21582174
RETURN_IF_ERROR(
2159-
compiler_enter_scope(c, name, scope_type, (void *)s, firstlineno, NULL));
2175+
compiler_enter_scope(c, name, scope_type, (void *)s, firstlineno, NULL, &umd));
21602176

21612177
Py_ssize_t first_instr = 0;
21622178
PyObject *docstring = _PyAST_GetDocString(body);
@@ -2181,9 +2197,9 @@ compiler_function_body(struct compiler *c, stmt_ty s, int is_async, Py_ssize_t f
21812197
}
21822198
Py_CLEAR(docstring);
21832199

2184-
c->u->u_metadata.u_argcount = asdl_seq_LEN(args->args);
2185-
c->u->u_metadata.u_posonlyargcount = asdl_seq_LEN(args->posonlyargs);
2186-
c->u->u_metadata.u_kwonlyargcount = asdl_seq_LEN(args->kwonlyargs);
2200+
assert(c->u->u_metadata.u_argcount == asdl_seq_LEN(args->args));
2201+
assert(c->u->u_metadata.u_posonlyargcount == asdl_seq_LEN(args->posonlyargs));
2202+
assert(c->u->u_metadata.u_kwonlyargcount == asdl_seq_LEN(args->kwonlyargs));
21872203

21882204
NEW_JUMP_TARGET_LABEL(c, start);
21892205
USE_LABEL(c, start);
@@ -2222,7 +2238,7 @@ compiler_function_body(struct compiler *c, stmt_ty s, int is_async, Py_ssize_t f
22222238
}
22232239

22242240
static int
2225-
compiler_function(struct compiler *c, stmt_ty s, int is_async)
2241+
codegen_function(struct compiler *c, stmt_ty s, int is_async)
22262242
{
22272243
arguments_ty args;
22282244
expr_ty returns;
@@ -2282,8 +2298,11 @@ compiler_function(struct compiler *c, stmt_ty s, int is_async)
22822298
if (!type_params_name) {
22832299
return ERROR;
22842300
}
2301+
_PyCompile_CodeUnitMetadata umd = {
2302+
.u_argcount = num_typeparam_args,
2303+
};
22852304
if (compiler_enter_scope(c, type_params_name, COMPILER_SCOPE_ANNOTATIONS,
2286-
(void *)type_params, firstlineno, NULL) == -1) {
2305+
(void *)type_params, firstlineno, NULL, &umd) == -1) {
22872306
Py_DECREF(type_params_name);
22882307
return ERROR;
22892308
}
@@ -2294,7 +2313,7 @@ compiler_function(struct compiler *c, stmt_ty s, int is_async)
22942313
}
22952314
}
22962315

2297-
int annotations_flag = compiler_annotations(c, loc, args, returns);
2316+
int annotations_flag = codegen_annotations(c, loc, args, returns);
22982317
if (annotations_flag < 0) {
22992318
if (is_generic) {
23002319
compiler_exit_scope(c);
@@ -2314,7 +2333,7 @@ compiler_function(struct compiler *c, stmt_ty s, int is_async)
23142333
ADDOP_I_IN_SCOPE(c, loc, SWAP, 2);
23152334
ADDOP_I_IN_SCOPE(c, loc, CALL_INTRINSIC_2, INTRINSIC_SET_FUNCTION_TYPE_PARAMS);
23162335

2317-
c->u->u_metadata.u_argcount = num_typeparam_args;
2336+
assert(c->u->u_metadata.u_argcount == num_typeparam_args);
23182337
PyCodeObject *co = optimize_and_assemble(c, 0);
23192338
compiler_exit_scope(c);
23202339
if (co == NULL) {
@@ -2365,7 +2384,7 @@ compiler_class_body(struct compiler *c, stmt_ty s, int firstlineno)
23652384
/* 1. compile the class body into a code object */
23662385
RETURN_IF_ERROR(
23672386
compiler_enter_scope(c, s->v.ClassDef.name, COMPILER_SCOPE_CLASS,
2368-
(void *)s, firstlineno, s->v.ClassDef.name));
2387+
(void *)s, firstlineno, s->v.ClassDef.name, NULL));
23692388

23702389
location loc = LOCATION(firstlineno, firstlineno, 0, 0);
23712390
/* load (global) __name__ ... */
@@ -2508,7 +2527,7 @@ compiler_class(struct compiler *c, stmt_ty s)
25082527
return ERROR;
25092528
}
25102529
if (compiler_enter_scope(c, type_params_name, COMPILER_SCOPE_ANNOTATIONS,
2511-
(void *)type_params, firstlineno, s->v.ClassDef.name) == -1) {
2530+
(void *)type_params, firstlineno, s->v.ClassDef.name, NULL) == -1) {
25122531
Py_DECREF(type_params_name);
25132532
return ERROR;
25142533
}
@@ -2592,7 +2611,7 @@ compiler_typealias_body(struct compiler *c, stmt_ty s)
25922611
PyObject *defaults = PyTuple_Pack(1, _PyLong_GetOne());
25932612
ADDOP_LOAD_CONST_NEW(c, loc, defaults);
25942613
RETURN_IF_ERROR(
2595-
compiler_setup_annotations_scope(c, LOC(s), s, name));
2614+
codegen_setup_annotations_scope(c, LOC(s), s, name));
25962615
/* Make None the first constant, so the evaluate function can't have a
25972616
docstring. */
25982617
RETURN_IF_ERROR(compiler_add_const(c, Py_None));
@@ -2627,7 +2646,7 @@ compiler_typealias(struct compiler *c, stmt_ty s)
26272646
return ERROR;
26282647
}
26292648
if (compiler_enter_scope(c, type_params_name, COMPILER_SCOPE_ANNOTATIONS,
2630-
(void *)type_params, loc.lineno, NULL) == -1) {
2649+
(void *)type_params, loc.lineno, NULL, NULL) == -1) {
26312650
Py_DECREF(type_params_name);
26322651
return ERROR;
26332652
}
@@ -2889,18 +2908,23 @@ compiler_lambda(struct compiler *c, expr_ty e)
28892908
return ERROR;
28902909
}
28912910

2911+
_PyCompile_CodeUnitMetadata umd = {
2912+
.u_argcount = asdl_seq_LEN(args->args),
2913+
.u_posonlyargcount = asdl_seq_LEN(args->posonlyargs),
2914+
.u_kwonlyargcount = asdl_seq_LEN(args->kwonlyargs),
2915+
};
28922916
_Py_DECLARE_STR(anon_lambda, "<lambda>");
28932917
RETURN_IF_ERROR(
28942918
compiler_enter_scope(c, &_Py_STR(anon_lambda), COMPILER_SCOPE_LAMBDA,
2895-
(void *)e, e->lineno, NULL));
2919+
(void *)e, e->lineno, NULL, &umd));
28962920

28972921
/* Make None the first constant, so the lambda can't have a
28982922
docstring. */
28992923
RETURN_IF_ERROR(compiler_add_const(c, Py_None));
29002924

2901-
c->u->u_metadata.u_argcount = asdl_seq_LEN(args->args);
2902-
c->u->u_metadata.u_posonlyargcount = asdl_seq_LEN(args->posonlyargs);
2903-
c->u->u_metadata.u_kwonlyargcount = asdl_seq_LEN(args->kwonlyargs);
2925+
assert(c->u->u_metadata.u_argcount == asdl_seq_LEN(args->args));
2926+
assert(c->u->u_metadata.u_posonlyargcount == asdl_seq_LEN(args->posonlyargs));
2927+
assert(c->u->u_metadata.u_kwonlyargcount == asdl_seq_LEN(args->kwonlyargs));
29042928
VISIT_IN_SCOPE(c, expr, e->v.Lambda.body);
29052929
if (SYMTABLE_ENTRY(c)->ste_generator) {
29062930
co = optimize_and_assemble(c, 0);
@@ -3867,7 +3891,7 @@ compiler_visit_stmt(struct compiler *c, stmt_ty s)
38673891

38683892
switch (s->kind) {
38693893
case FunctionDef_kind:
3870-
return compiler_function(c, s, 0);
3894+
return codegen_function(c, s, 0);
38713895
case ClassDef_kind:
38723896
return compiler_class(c, s);
38733897
case TypeAlias_kind:
@@ -3949,7 +3973,7 @@ compiler_visit_stmt(struct compiler *c, stmt_ty s)
39493973
case With_kind:
39503974
return compiler_with(c, s, 0);
39513975
case AsyncFunctionDef_kind:
3952-
return compiler_function(c, s, 1);
3976+
return codegen_function(c, s, 1);
39533977
case AsyncWith_kind:
39543978
return compiler_async_with(c, s, 0);
39553979
case AsyncFor_kind:
@@ -5674,7 +5698,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
56745698
}
56755699
else {
56765700
if (compiler_enter_scope(c, name, COMPILER_SCOPE_COMPREHENSION,
5677-
(void *)e, e->lineno, NULL) < 0) {
5701+
(void *)e, e->lineno, NULL, NULL) < 0) {
56785702
goto error;
56795703
}
56805704
}

0 commit comments

Comments
 (0)