python · ericsnowcurrently · May 5, 2025 · May 1, 2025 · May 1, 2025 · May 1, 2025
diff --git a/Include/cpython/funcobject.h b/Include/cpython/funcobject.h
@@ -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;
 }

diff --git a/Include/internal/pycore_code.h b/Include/internal/pycore_code.h
@@ -565,6 +565,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 *);
 
 

diff --git a/Lib/test/test_code.py b/Lib/test/test_code.py
@@ -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])