Skip to content

gh-132775: Add _PyCode_GetVarCounts() #133128

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Apr 30, 2025
Merged

gh-132775: Add _PyCode_GetVarCounts() #133128

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
d95aef5
Add _PyCode_GetVarCounts().
ericsnowcurrently Apr 14, 2025
aaa98ac
Fix the tests.
ericsnowcurrently Apr 29, 2025
c8181cc
Add PyFunction_GET_BUILTINS().
ericsnowcurrently Apr 29, 2025
ddbe9bc
Optionally identify globals/builtins counts.
ericsnowcurrently Apr 29, 2025
195a1ec
Pass tstate to _PyCode_SetUnboundVarCounts().
ericsnowcurrently Apr 29, 2025
998aeed
Add some comments.
ericsnowcurrently Apr 30, 2025
a727d7b
Fix a compiler warning.
ericsnowcurrently Apr 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions Include/cpython/funcobject.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,11 @@ static inline PyObject* PyFunction_GET_GLOBALS(PyObject *func) {
}
#define PyFunction_GET_GLOBALS(func) PyFunction_GET_GLOBALS(_PyObject_CAST(func))

static inline PyObject* PyFunction_GET_BUILTINS(PyObject *func) {
return _PyFunction_CAST(func)->func_builtins;
}
#define PyFunction_GET_BUILTINS(func) PyFunction_GET_BUILTINS(_PyObject_CAST(func))

static inline PyObject* PyFunction_GET_MODULE(PyObject *func) {
return _PyFunction_CAST(func)->func_module;
}
Expand Down
51 changes: 51 additions & 0 deletions Include/internal/pycore_code.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -564,6 +564,57 @@ extern int _Py_ClearUnusedTLBC(PyInterpreterState *interp);
#endif


typedef struct {
int total;
struct co_locals_counts {
int total;
struct {
int total;
int numposonly;
int numposorkw;
int numkwonly;
int varargs;
int varkwargs;
} args;
int numpure;
struct {
int total;
// numargs does not contribute to locals.total.
int numargs;
int numothers;
} cells;
struct {
int total;
int numpure;
int numcells;
} hidden;
} locals;
int numfree; // nonlocal
struct co_unbound_counts {
int total;
struct {
int total;
int numglobal;
int numbuiltin;
int numunknown;
} globals;
int numattrs;
int numunknown;
} unbound;
} _PyCode_var_counts_t;

PyAPI_FUNC(void) _PyCode_GetVarCounts(
PyCodeObject *,
_PyCode_var_counts_t *);
PyAPI_FUNC(int) _PyCode_SetUnboundVarCounts(
PyThreadState *,
PyCodeObject *,
_PyCode_var_counts_t *,
PyObject *globalnames,
PyObject *attrnames,
PyObject *globalsns,
PyObject *builtinsns);

PyAPI_FUNC(int) _PyCode_ReturnsOnlyNone(PyCodeObject *);


Expand Down
230 changes: 230 additions & 0 deletions Lib/test/test_code.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -777,6 +777,236 @@ def test_local_kinds(self):
kinds = _testinternalcapi.get_co_localskinds(func.__code__)
self.assertEqual(kinds, expected)

@unittest.skipIf(_testinternalcapi is None, "missing _testinternalcapi")
def test_var_counts(self):
self.maxDiff = None
def new_var_counts(*,
posonly=0,
posorkw=0,
kwonly=0,
varargs=0,
varkwargs=0,
purelocals=0,
argcells=0,
othercells=0,
freevars=0,
globalvars=0,
attrs=0,
unknown=0,
):
nargvars = posonly + posorkw + kwonly + varargs + varkwargs
nlocals = nargvars + purelocals + othercells
if isinstance(globalvars, int):
globalvars = {
'total': globalvars,
'numglobal': 0,
'numbuiltin': 0,
'numunknown': globalvars,
}
else:
g_numunknown = 0
if isinstance(globalvars, dict):
numglobal = globalvars['numglobal']
numbuiltin = globalvars['numbuiltin']
size = 2
if 'numunknown' in globalvars:
g_numunknown = globalvars['numunknown']
size += 1
assert len(globalvars) == size, globalvars
else:
assert not isinstance(globalvars, str), repr(globalvars)
try:
numglobal, numbuiltin = globalvars
except ValueError:
numglobal, numbuiltin, g_numunknown = globalvars
globalvars = {
'total': numglobal + numbuiltin + g_numunknown,
'numglobal': numglobal,
'numbuiltin': numbuiltin,
'numunknown': g_numunknown,
}
unbound = globalvars['total'] + attrs + unknown
return {
'total': nlocals + freevars + unbound,
'locals': {
'total': nlocals,
'args': {
'total': nargvars,
'numposonly': posonly,
'numposorkw': posorkw,
'numkwonly': kwonly,
'varargs': varargs,
'varkwargs': varkwargs,
},
'numpure': purelocals,
'cells': {
'total': argcells + othercells,
'numargs': argcells,
'numothers': othercells,
},
'hidden': {
'total': 0,
'numpure': 0,
'numcells': 0,
},
},
'numfree': freevars,
'unbound': {
'total': unbound,
'globals': globalvars,
'numattrs': attrs,
'numunknown': unknown,
},
}

import test._code_definitions as defs
funcs = {
defs.spam_minimal: new_var_counts(),
defs.spam_full: new_var_counts(
posonly=2,
posorkw=2,
kwonly=2,
varargs=1,
varkwargs=1,
purelocals=4,
globalvars=3,
attrs=1,
),
defs.spam: new_var_counts(
posorkw=1,
),
defs.spam_N: new_var_counts(
posorkw=1,
purelocals=1,
),
defs.spam_C: new_var_counts(
posorkw=1,
purelocals=1,
argcells=1,
othercells=1,
),
defs.spam_NN: new_var_counts(
posorkw=1,
purelocals=1,
),
defs.spam_NC: new_var_counts(
posorkw=1,
purelocals=1,
argcells=1,
othercells=1,
),
defs.spam_CN: new_var_counts(
posorkw=1,
purelocals=1,
argcells=1,
othercells=1,
),
defs.spam_CC: new_var_counts(
posorkw=1,
purelocals=1,
argcells=1,
othercells=1,
),
defs.eggs_nested: new_var_counts(
posorkw=1,
),
defs.eggs_closure: new_var_counts(
posorkw=1,
freevars=2,
),
defs.eggs_nested_N: new_var_counts(
posorkw=1,
purelocals=1,
),
defs.eggs_nested_C: new_var_counts(
posorkw=1,
purelocals=1,
argcells=1,
freevars=2,
),
defs.eggs_closure_N: new_var_counts(
posorkw=1,
purelocals=1,
freevars=2,
),
defs.eggs_closure_C: new_var_counts(
posorkw=1,
purelocals=1,
argcells=1,
othercells=1,
freevars=2,
),
defs.ham_nested: new_var_counts(
posorkw=1,
),
defs.ham_closure: new_var_counts(
posorkw=1,
freevars=3,
),
defs.ham_C_nested: new_var_counts(
posorkw=1,
),
defs.ham_C_closure: new_var_counts(
posorkw=1,
freevars=4,
),
}
assert len(funcs) == len(defs.FUNCTIONS), (len(funcs), len(defs.FUNCTIONS))
for func in defs.FUNCTIONS:
with self.subTest(func):
expected = funcs[func]
counts = _testinternalcapi.get_code_var_counts(func.__code__)
self.assertEqual(counts, expected)

def func_with_globals_and_builtins():
mod1 = _testinternalcapi
mod2 = dis
mods = (mod1, mod2)
checks = tuple(callable(m) for m in mods)
return callable(mod2), tuple(mods), list(mods), checks

func = func_with_globals_and_builtins
with self.subTest(f'{func} code'):
expected = new_var_counts(
purelocals=4,
globalvars=5,
)
counts = _testinternalcapi.get_code_var_counts(func.__code__)
self.assertEqual(counts, expected)

with self.subTest(f'{func} with own globals and builtins'):
expected = new_var_counts(
purelocals=4,
globalvars=(2, 3),
)
counts = _testinternalcapi.get_code_var_counts(func)
self.assertEqual(counts, expected)

with self.subTest(f'{func} without globals'):
expected = new_var_counts(
purelocals=4,
globalvars=(0, 3, 2),
)
counts = _testinternalcapi.get_code_var_counts(func, globalsns={})
self.assertEqual(counts, expected)

with self.subTest(f'{func} without both'):
expected = new_var_counts(
purelocals=4,
globalvars=5,
)
counts = _testinternalcapi.get_code_var_counts(func, globalsns={},
builtinsns={})
self.assertEqual(counts, expected)

with self.subTest(f'{func} without builtins'):
expected = new_var_counts(
purelocals=4,
globalvars=(2, 0, 3),
)
counts = _testinternalcapi.get_code_var_counts(func, builtinsns={})
self.assertEqual(counts, expected)


def isinterned(s):
return s is sys.intern(('_' + s + '_')[1:-1])
Expand Down
Loading
Loading