From 1ffdc1b9df998540664cf7c37a8ffe6c4e1a79b2 Mon Sep 17 00:00:00 2001 From: Eric Snow Date: Fri, 25 Apr 2025 15:42:44 -0600 Subject: [PATCH 1/3] Track code object local kinds for arguments. --- Include/internal/pycore_code.h | 14 ++++---- Python/assemble.c | 65 ++++++++++++++++++++++++++++------ 2 files changed, 62 insertions(+), 17 deletions(-) diff --git a/Include/internal/pycore_code.h b/Include/internal/pycore_code.h index 2839b9b7ebe0fb..c9a6fae006551f 100644 --- a/Include/internal/pycore_code.h +++ b/Include/internal/pycore_code.h @@ -177,12 +177,14 @@ typedef struct { */ // Note that these all fit within a byte, as do combinations. -// Later, we will use the smaller numbers to differentiate the different -// kinds of locals (e.g. pos-only arg, varkwargs, local-only). -#define CO_FAST_HIDDEN 0x10 -#define CO_FAST_LOCAL 0x20 -#define CO_FAST_CELL 0x40 -#define CO_FAST_FREE 0x80 +#define CO_FAST_ARG_POS (0x02) // pos-only, pos-or-kw, varargs +#define CO_FAST_ARG_KW (0x04) // kw-only, pos-or-kw, varkwargs +#define CO_FAST_ARG_VAR (0x08) // varargs, varkwargs +#define CO_FAST_ARG (CO_FAST_ARG_POS | CO_FAST_ARG_KW | CO_FAST_ARG_VAR) +#define CO_FAST_HIDDEN (0x10) +#define CO_FAST_LOCAL (0x20) +#define CO_FAST_CELL (0x40) +#define CO_FAST_FREE (0x80) typedef unsigned char _PyLocals_Kind; diff --git a/Python/assemble.c b/Python/assemble.c index 5f7af682eca0ef..4099c40aa7461b 100644 --- a/Python/assemble.c +++ b/Python/assemble.c @@ -482,28 +482,62 @@ extern void _Py_set_localsplus_info(int, PyObject *, unsigned char, static int compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus, - PyObject *names, PyObject *kinds) + int flags, PyObject *names, PyObject *kinds) { + // Compute the arg flags. + // Currently arg vars fill the first portion of the list. + char *argkinds = PyMem_Calloc(nlocalsplus, sizeof(char)); + if (argkinds == NULL) { + PyErr_NoMemory(); + return ERROR; + } + int numargvars = 0; + int max = (int)umd->u_posonlyargcount; + for (; numargvars < max; numargvars++) { + argkinds[numargvars] = CO_FAST_ARG_POS; + } + max += (int)umd->u_argcount; + for (; numargvars < max; numargvars++) { + argkinds[numargvars] = CO_FAST_ARG_POS | CO_FAST_ARG_KW; + } + max += (int)umd->u_kwonlyargcount; + for (; numargvars < max; numargvars++) { + argkinds[numargvars] = CO_FAST_ARG_KW; + } + if (flags & CO_VARARGS) { + argkinds[numargvars] = CO_FAST_ARG_VAR | CO_FAST_ARG_POS; + numargvars += 1; + } + if (flags & CO_VARKEYWORDS) { + argkinds[numargvars] = CO_FAST_ARG_VAR | CO_FAST_ARG_KW; + numargvars += 1; + } + + // Set the locals kinds. PyObject *k, *v; Py_ssize_t pos = 0; while (PyDict_Next(umd->u_varnames, &pos, &k, &v)) { int offset = PyLong_AsInt(v); if (offset == -1 && PyErr_Occurred()) { - return ERROR; + goto error; } assert(offset >= 0); assert(offset < nlocalsplus); - // For now we do not distinguish arg kinds. - _PyLocals_Kind kind = CO_FAST_LOCAL; + _PyLocals_Kind kind = CO_FAST_LOCAL | argkinds[offset]; + int has_key = PyDict_Contains(umd->u_fasthidden, k); - RETURN_IF_ERROR(has_key); + if (has_key < 0) { + goto error; + } if (has_key) { kind |= CO_FAST_HIDDEN; } has_key = PyDict_Contains(umd->u_cellvars, k); - RETURN_IF_ERROR(has_key); + if (has_key < 0) { + goto error; + } if (has_key) { kind |= CO_FAST_CELL; } @@ -517,7 +551,9 @@ compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus, pos = 0; while (PyDict_Next(umd->u_cellvars, &pos, &k, &v)) { int has_name = PyDict_Contains(umd->u_varnames, k); - RETURN_IF_ERROR(has_name); + if (has_name < 0) { + goto error; + } if (has_name) { // Skip cells that are already covered by locals. numdropped += 1; @@ -526,7 +562,7 @@ compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus, cellvar_offset = PyLong_AsInt(v); if (cellvar_offset == -1 && PyErr_Occurred()) { - return ERROR; + goto error; } assert(cellvar_offset >= 0); cellvar_offset += nlocals - numdropped; @@ -538,7 +574,7 @@ compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus, while (PyDict_Next(umd->u_freevars, &pos, &k, &v)) { int offset = PyLong_AsInt(v); if (offset == -1 && PyErr_Occurred()) { - return ERROR; + goto error; } assert(offset >= 0); offset += nlocals - numdropped; @@ -549,7 +585,12 @@ compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus, assert(offset > cellvar_offset); _Py_set_localsplus_info(offset, k, CO_FAST_FREE, names, kinds); } + PyMem_Free(argkinds); return SUCCESS; + +error: + PyMem_Free(argkinds); + return ERROR; } static PyCodeObject * @@ -594,8 +635,10 @@ makecode(_PyCompile_CodeUnitMetadata *umd, struct assembler *a, PyObject *const_ if (localspluskinds == NULL) { goto error; } - if (compute_localsplus_info(umd, nlocalsplus, - localsplusnames, localspluskinds) == ERROR) { + if (compute_localsplus_info( + umd, nlocalsplus, code_flags, + localsplusnames, localspluskinds) == ERROR) + { goto error; } From 6d22d3b2790579f42d3e9aa8a0888ed4c63b381a Mon Sep 17 00:00:00 2001 From: Eric Snow Date: Mon, 28 Apr 2025 16:43:43 -0600 Subject: [PATCH 2/3] Avoid allocating temporary memory. --- Python/assemble.c | 104 ++++++++++++++++++---------------------------- 1 file changed, 41 insertions(+), 63 deletions(-) diff --git a/Python/assemble.c b/Python/assemble.c index 4099c40aa7461b..8cc2d50a3227f8 100644 --- a/Python/assemble.c +++ b/Python/assemble.c @@ -484,65 +484,50 @@ static int compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus, int flags, PyObject *names, PyObject *kinds) { - // Compute the arg flags. - // Currently arg vars fill the first portion of the list. - char *argkinds = PyMem_Calloc(nlocalsplus, sizeof(char)); - if (argkinds == NULL) { - PyErr_NoMemory(); - return ERROR; - } - int numargvars = 0; - int max = (int)umd->u_posonlyargcount; - for (; numargvars < max; numargvars++) { - argkinds[numargvars] = CO_FAST_ARG_POS; - } - max += (int)umd->u_argcount; - for (; numargvars < max; numargvars++) { - argkinds[numargvars] = CO_FAST_ARG_POS | CO_FAST_ARG_KW; - } - max += (int)umd->u_kwonlyargcount; - for (; numargvars < max; numargvars++) { - argkinds[numargvars] = CO_FAST_ARG_KW; - } - if (flags & CO_VARARGS) { - argkinds[numargvars] = CO_FAST_ARG_VAR | CO_FAST_ARG_POS; - numargvars += 1; - } - if (flags & CO_VARKEYWORDS) { - argkinds[numargvars] = CO_FAST_ARG_VAR | CO_FAST_ARG_KW; - numargvars += 1; - } - - // Set the locals kinds. PyObject *k, *v; Py_ssize_t pos = 0; - while (PyDict_Next(umd->u_varnames, &pos, &k, &v)) { - int offset = PyLong_AsInt(v); - if (offset == -1 && PyErr_Occurred()) { - goto error; - } - assert(offset >= 0); - assert(offset < nlocalsplus); - _PyLocals_Kind kind = CO_FAST_LOCAL | argkinds[offset]; + // Set the locals kinds. Arg vars fill the first portion of the list. + struct { + int count; + _PyLocals_Kind kind; + } argvarkinds[6] = { + {(int)umd->u_posonlyargcount, CO_FAST_ARG_POS}, + {(int)umd->u_argcount, CO_FAST_ARG_POS | CO_FAST_ARG_KW}, + {(int)umd->u_kwonlyargcount, CO_FAST_ARG_KW}, + {!!(flags & CO_VARARGS), CO_FAST_ARG_VAR | CO_FAST_ARG_POS}, + {!!(flags & CO_VARKEYWORDS), CO_FAST_ARG_VAR | CO_FAST_ARG_KW}, + {-1, 0}, // the remaining local vars + }; + int max = 0; + for (int i = 0; i < 6; i++) { + max = argvarkinds[i].count < 0 + ? INT_MAX + : max + argvarkinds[i].count; + while (pos < max && PyDict_Next(umd->u_varnames, &pos, &k, &v)) { + int offset = PyLong_AsInt(v); + if (offset == -1 && PyErr_Occurred()) { + return ERROR; + } + assert(offset >= 0); + assert(offset < nlocalsplus); + + _PyLocals_Kind kind = CO_FAST_LOCAL | argvarkinds[i].kind; - int has_key = PyDict_Contains(umd->u_fasthidden, k); - if (has_key < 0) { - goto error; - } - if (has_key) { - kind |= CO_FAST_HIDDEN; - } + int has_key = PyDict_Contains(umd->u_fasthidden, k); + RETURN_IF_ERROR(has_key); + if (has_key) { + kind |= CO_FAST_HIDDEN; + } - has_key = PyDict_Contains(umd->u_cellvars, k); - if (has_key < 0) { - goto error; - } - if (has_key) { - kind |= CO_FAST_CELL; - } + has_key = PyDict_Contains(umd->u_cellvars, k); + RETURN_IF_ERROR(has_key); + if (has_key) { + kind |= CO_FAST_CELL; + } - _Py_set_localsplus_info(offset, k, kind, names, kinds); + _Py_set_localsplus_info(offset, k, kind, names, kinds); + } } int nlocals = (int)PyDict_GET_SIZE(umd->u_varnames); @@ -551,9 +536,7 @@ compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus, pos = 0; while (PyDict_Next(umd->u_cellvars, &pos, &k, &v)) { int has_name = PyDict_Contains(umd->u_varnames, k); - if (has_name < 0) { - goto error; - } + RETURN_IF_ERROR(has_name); if (has_name) { // Skip cells that are already covered by locals. numdropped += 1; @@ -562,7 +545,7 @@ compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus, cellvar_offset = PyLong_AsInt(v); if (cellvar_offset == -1 && PyErr_Occurred()) { - goto error; + return ERROR; } assert(cellvar_offset >= 0); cellvar_offset += nlocals - numdropped; @@ -574,7 +557,7 @@ compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus, while (PyDict_Next(umd->u_freevars, &pos, &k, &v)) { int offset = PyLong_AsInt(v); if (offset == -1 && PyErr_Occurred()) { - goto error; + return ERROR; } assert(offset >= 0); offset += nlocals - numdropped; @@ -585,12 +568,7 @@ compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus, assert(offset > cellvar_offset); _Py_set_localsplus_info(offset, k, CO_FAST_FREE, names, kinds); } - PyMem_Free(argkinds); return SUCCESS; - -error: - PyMem_Free(argkinds); - return ERROR; } static PyCodeObject * From 144b98d93da7d6fcfc66ff0ae4e82348f665418d Mon Sep 17 00:00:00 2001 From: Eric Snow Date: Mon, 28 Apr 2025 16:44:01 -0600 Subject: [PATCH 3/3] Add tests. --- Lib/test/_code_definitions.py | 163 ++++++++++++++++++++++++++ Lib/test/_crossinterp_definitions.py | 167 +-------------------------- Lib/test/test_code.py | 127 ++++++++++++++++++++ Modules/_testinternalcapi.c | 32 +++++ 4 files changed, 324 insertions(+), 165 deletions(-) create mode 100644 Lib/test/_code_definitions.py diff --git a/Lib/test/_code_definitions.py b/Lib/test/_code_definitions.py new file mode 100644 index 00000000000000..06cf6a10231f51 --- /dev/null +++ b/Lib/test/_code_definitions.py @@ -0,0 +1,163 @@ + + +def spam_minimal(): + # no arg defaults or kwarg defaults + # no annotations + # no local vars + # no free vars + # no globals + # no builtins + # no attr access (names) + # no code + return + + +def spam_full(a, b, /, c, d:int=1, *args, e, f:object=None, **kwargs) -> tuple: + # arg defaults, kwarg defaults + # annotations + # all kinds of local vars, except cells + # no free vars + # some globals + # some builtins + # some attr access (names) + x = args + y = kwargs + z = (a, b, c, d) + kwargs['e'] = e + kwargs['f'] = f + extras = list((x, y, z, spam, spam.__name__)) + return tuple(a, b, c, d, e, f, args, kwargs), extras + + +def spam(x): + return x, None + + +def spam_N(x): + def eggs_nested(y): + return None, y + return eggs_nested, x + + +def spam_C(x): + a = 1 + def eggs_closure(y): + return None, y, a, x + return eggs_closure, a, x + + +def spam_NN(x): + def eggs_nested_N(y): + def ham_nested(z): + return None, z + return ham_nested, y + return eggs_nested_N, x + + +def spam_NC(x): + a = 1 + def eggs_nested_C(y): + def ham_closure(z): + return None, z, y, a, x + return ham_closure, y + return eggs_nested_C, a, x + + +def spam_CN(x): + a = 1 + def eggs_closure_N(y): + def ham_C_nested(z): + return None, z + return ham_C_nested, y, a, x + return eggs_closure_N, a, x + + +def spam_CC(x): + a = 1 + def eggs_closure_C(y): + b = 2 + def ham_C_closure(z): + return None, z, b, y, a, x + return ham_C_closure, b, y, a, x + return eggs_closure_C, a, x + + +eggs_nested, *_ = spam_N(1) +eggs_closure, *_ = spam_C(1) +eggs_nested_N, *_ = spam_NN(1) +eggs_nested_C, *_ = spam_NC(1) +eggs_closure_N, *_ = spam_CN(1) +eggs_closure_C, *_ = spam_CC(1) + +ham_nested, *_ = eggs_nested_N(2) +ham_closure, *_ = eggs_nested_C(2) +ham_C_nested, *_ = eggs_closure_N(2) +ham_C_closure, *_ = eggs_closure_C(2) + + +TOP_FUNCTIONS = [ + # shallow + spam_minimal, + spam_full, + spam, + # outer func + spam_N, + spam_C, + spam_NN, + spam_NC, + spam_CN, + spam_CC, +] +NESTED_FUNCTIONS = [ + # inner func + eggs_nested, + eggs_closure, + eggs_nested_N, + eggs_nested_C, + eggs_closure_N, + eggs_closure_C, + # inner inner func + ham_nested, + ham_closure, + ham_C_nested, + ham_C_closure, +] +FUNCTIONS = [ + *TOP_FUNCTIONS, + *NESTED_FUNCTIONS, +] + + +# generators + +def gen_spam_1(*args): + for arg in args: + yield arg + + +def gen_spam_2(*args): + yield from args + + +async def async_spam(): + pass +coro_spam = async_spam() +coro_spam.close() + + +async def asyncgen_spam(*args): + for arg in args: + yield arg +asynccoro_spam = asyncgen_spam(1, 2, 3) + + +FUNCTION_LIKE = [ + gen_spam_1, + gen_spam_2, + async_spam, + asyncgen_spam, +] +FUNCTION_LIKE_APPLIED = [ + coro_spam, # actually FunctionType? + asynccoro_spam, # actually FunctionType? +] diff --git a/Lib/test/_crossinterp_definitions.py b/Lib/test/_crossinterp_definitions.py index 0d5f6c7db064d3..5c134427825c7d 100644 --- a/Lib/test/_crossinterp_definitions.py +++ b/Lib/test/_crossinterp_definitions.py @@ -3,172 +3,9 @@ ####################################### -# functions - -def spam_minimal(): - # no arg defaults or kwarg defaults - # no annotations - # no local vars - # no free vars - # no globals - # no builtins - # no attr access (names) - # no code - return - - -def spam_full(a, b, /, c, d:int=1, *args, e, f:object=None, **kwargs) -> tuple: - # arg defaults, kwarg defaults - # annotations - # all kinds of local vars, except cells - # no free vars - # some globals - # some builtins - # some attr access (names) - x = args - y = kwargs - z = (a, b, c, d) - kwargs['e'] = e - kwargs['f'] = f - extras = list((x, y, z, spam, spam.__name__)) - return tuple(a, b, c, d, e, f, args, kwargs), extras - - -def spam(x): - return x, None - - -def spam_N(x): - def eggs_nested(y): - return None, y - return eggs_nested, x - - -def spam_C(x): - a = 1 - def eggs_closure(y): - return None, y, a, x - return eggs_closure, a, x - - -def spam_NN(x): - def eggs_nested_N(y): - def ham_nested(z): - return None, z - return ham_nested, y - return eggs_nested_N, x - - -def spam_NC(x): - a = 1 - def eggs_nested_C(y): - def ham_closure(z): - return None, z, y, a, x - return ham_closure, y - return eggs_nested_C, a, x - - -def spam_CN(x): - a = 1 - def eggs_closure_N(y): - def ham_C_nested(z): - return None, z - return ham_C_nested, y, a, x - return eggs_closure_N, a, x - - -def spam_CC(x): - a = 1 - def eggs_closure_C(y): - b = 2 - def ham_C_closure(z): - return None, z, b, y, a, x - return ham_C_closure, b, y, a, x - return eggs_closure_N, a, x - - -eggs_nested, *_ = spam_N(1) -eggs_closure, *_ = spam_C(1) -eggs_nested_N, *_ = spam_NN(1) -eggs_nested_C, *_ = spam_NC(1) -eggs_closure_N, *_ = spam_CN(1) -eggs_closure_C, *_ = spam_CC(1) - -ham_nested, *_ = eggs_nested_N(2) -ham_closure, *_ = eggs_nested_C(2) -ham_C_nested, *_ = eggs_closure_N(2) -ham_C_closure, *_ = eggs_closure_C(2) - - -TOP_FUNCTIONS = [ - # shallow - spam_minimal, - spam_full, - spam, - # outer func - spam_N, - spam_C, - spam_NN, - spam_NC, - spam_CN, - spam_CC, -] -NESTED_FUNCTIONS = [ - # inner func - eggs_nested, - eggs_closure, - eggs_nested_N, - eggs_nested_C, - eggs_closure_N, - eggs_closure_C, - # inner inner func - ham_nested, - ham_closure, - ham_C_nested, - ham_C_closure, -] -FUNCTIONS = [ - *TOP_FUNCTIONS, - *NESTED_FUNCTIONS, -] - - -####################################### -# function-like - -# generators - -def gen_spam_1(*args): - for arg in args: - yield arg - +# functions and generators -def gen_spam_2(*args): - yield from args - - -async def async_spam(): - pass -coro_spam = async_spam() -coro_spam.close() - - -async def asyncgen_spam(*args): - for arg in args: - yield arg -asynccoro_spam = asyncgen_spam(1, 2, 3) - - -FUNCTION_LIKE = [ - gen_spam_1, - gen_spam_2, - async_spam, - asyncgen_spam, -] -FUNCTION_LIKE_APPLIED = [ - coro_spam, # actually FunctionType? - asynccoro_spam, # actually FunctionType? -] +from test._code_definitions import * ####################################### diff --git a/Lib/test/test_code.py b/Lib/test/test_code.py index 2f459a46b5ad70..a811c67f0d0a30 100644 --- a/Lib/test/test_code.py +++ b/Lib/test/test_code.py @@ -216,6 +216,10 @@ from test.support.bytecode_helper import instructions_with_positions from opcode import opmap, opname from _testcapi import code_offset_to_line +try: + import _testinternalcapi +except ModuleNotFoundError: + _testinternalcapi = None COPY_FREE_VARS = opmap['COPY_FREE_VARS'] @@ -242,6 +246,7 @@ def dump(co): def external_getitem(self, i): return f"Foreign getitem: {super().__getitem__(i)}" + class CodeTest(unittest.TestCase): @cpython_only @@ -595,6 +600,128 @@ def test_code_equal_with_instrumentation(self): self.assertNotEqual(code1, code2) sys.settrace(None) + @unittest.skipIf(_testinternalcapi is None, "missing _testinternalcapi") + def test_local_kinds(self): + CO_FAST_ARG_POS = 0x02 + CO_FAST_ARG_KW = 0x04 + CO_FAST_ARG_VAR = 0x08 + CO_FAST_HIDDEN = 0x10 + CO_FAST_LOCAL = 0x20 + CO_FAST_CELL = 0x40 + CO_FAST_FREE = 0x80 + + POSONLY = CO_FAST_LOCAL | CO_FAST_ARG_POS + POSORKW = CO_FAST_LOCAL | CO_FAST_ARG_POS | CO_FAST_ARG_KW + KWONLY = CO_FAST_LOCAL | CO_FAST_ARG_KW + VARARGS = CO_FAST_LOCAL | CO_FAST_ARG_VAR | CO_FAST_ARG_POS + VARKWARGS = CO_FAST_LOCAL | CO_FAST_ARG_VAR | CO_FAST_ARG_KW + + import test._code_definitions as defs + funcs = { + defs.spam_minimal: {}, + defs.spam_full: { + 'a': POSONLY, + 'b': POSONLY, + 'c': POSORKW, + 'd': POSORKW, + 'e': KWONLY, + 'f': KWONLY, + 'args': VARARGS, + 'kwargs': VARKWARGS, + 'x': CO_FAST_LOCAL, + 'y': CO_FAST_LOCAL, + 'z': CO_FAST_LOCAL, + 'extras': CO_FAST_LOCAL, + }, + defs.spam: { + 'x': POSORKW, + }, + defs.spam_N: { + 'x': POSORKW, + 'eggs_nested': CO_FAST_LOCAL, + }, + defs.spam_C: { + 'x': POSORKW | CO_FAST_CELL, + 'a': CO_FAST_CELL, + 'eggs_closure': CO_FAST_LOCAL, + }, + defs.spam_NN: { + 'x': POSORKW, + 'eggs_nested_N': CO_FAST_LOCAL, + }, + defs.spam_NC: { + 'x': POSORKW | CO_FAST_CELL, + 'a': CO_FAST_CELL, + 'eggs_nested_C': CO_FAST_LOCAL, + }, + defs.spam_CN: { + 'x': POSORKW | CO_FAST_CELL, + 'a': CO_FAST_CELL, + 'eggs_closure_N': CO_FAST_LOCAL, + }, + defs.spam_CC: { + 'x': POSORKW | CO_FAST_CELL, + 'a': CO_FAST_CELL, + 'eggs_closure_C': CO_FAST_LOCAL, + }, + defs.eggs_nested: { + 'y': POSORKW, + }, + defs.eggs_closure: { + 'y': POSORKW, + 'x': CO_FAST_FREE, + 'a': CO_FAST_FREE, + }, + defs.eggs_nested_N: { + 'y': POSORKW, + 'ham_nested': CO_FAST_LOCAL, + }, + defs.eggs_nested_C: { + 'y': POSORKW | CO_FAST_CELL, + 'x': CO_FAST_FREE, + 'a': CO_FAST_FREE, + 'ham_closure': CO_FAST_LOCAL, + }, + defs.eggs_closure_N: { + 'y': POSORKW, + 'x': CO_FAST_FREE, + 'a': CO_FAST_FREE, + 'ham_C_nested': CO_FAST_LOCAL, + }, + defs.eggs_closure_C: { + 'y': POSORKW | CO_FAST_CELL, + 'b': CO_FAST_CELL, + 'x': CO_FAST_FREE, + 'a': CO_FAST_FREE, + 'ham_C_closure': CO_FAST_LOCAL, + }, + defs.ham_nested: { + 'z': POSORKW, + }, + defs.ham_closure: { + 'z': POSORKW, + 'y': CO_FAST_FREE, + 'x': CO_FAST_FREE, + 'a': CO_FAST_FREE, + }, + defs.ham_C_nested: { + 'z': POSORKW, + }, + defs.ham_C_closure: { + 'z': POSORKW, + 'y': CO_FAST_FREE, + 'b': CO_FAST_FREE, + 'x': CO_FAST_FREE, + 'a': CO_FAST_FREE, + }, + } + assert len(funcs) == len(defs.FUNCTIONS) + for func in defs.FUNCTIONS: + with self.subTest(func): + expected = funcs[func] + kinds = _testinternalcapi.get_co_localskinds(func.__code__) + self.assertEqual(kinds, expected) + def isinterned(s): return s is sys.intern(('_' + s + '_')[1:-1]) diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c index 0ef064fe80d173..9ce01b7c7fdb78 100644 --- a/Modules/_testinternalcapi.c +++ b/Modules/_testinternalcapi.c @@ -956,6 +956,37 @@ get_co_framesize(PyObject *self, PyObject *arg) return PyLong_FromLong(code->co_framesize); } +static PyObject * +get_co_localskinds(PyObject *self, PyObject *arg) +{ + if (!PyCode_Check(arg)) { + PyErr_SetString(PyExc_TypeError, "argument must be a code object"); + return NULL; + } + PyCodeObject *co = (PyCodeObject *)arg; + + PyObject *kinds = PyDict_New(); + if (kinds == NULL) { + return NULL; + } + for (int offset = 0; offset < co->co_nlocalsplus; offset++) { + PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, offset); + _PyLocals_Kind k = _PyLocals_GetKind(co->co_localspluskinds, offset); + PyObject *kind = PyLong_FromLong(k); + if (kind == NULL) { + Py_DECREF(kinds); + return NULL; + } + int res = PyDict_SetItem(kinds, name, kind); + Py_DECREF(kind); + if (res < 0) { + Py_DECREF(kinds); + return NULL; + } + } + return kinds; +} + static PyObject * jit_enabled(PyObject *self, PyObject *arg) { @@ -2075,6 +2106,7 @@ static PyMethodDef module_functions[] = { {"iframe_getline", iframe_getline, METH_O, NULL}, {"iframe_getlasti", iframe_getlasti, METH_O, NULL}, {"get_co_framesize", get_co_framesize, METH_O, NULL}, + {"get_co_localskinds", get_co_localskinds, METH_O, NULL}, {"jit_enabled", jit_enabled, METH_NOARGS, NULL}, #ifdef _Py_TIER2 {"add_executor_dependency", add_executor_dependency, METH_VARARGS, NULL},