gh-132775: Add _PyCode_GetVarCounts() (gh-133128)

This helper is useful in a variety of ways, including in demonstrating how the different counts relate to one another.

It will be used in a later change to help identify if a function is "stateless", meaning it doesn't have any free vars or globals.

Note that a majority of this change is tests.

5 files changed, 689 insertions, 0 deletions

diff --git a/Include/cpython/funcobject.h b/Include/cpython/funcobject.h
index 598cd33..18249b9 100644
--- 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
index 635d2b2..9b02e29 100644
--- 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
index 7cf09ee..1b6dfe7 100644
--- a/Lib/test/test_code.py
+++ b/Lib/test/test_code.py
@@ -777,6 +777,236 @@ class CodeTest(unittest.TestCase):
                 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])
diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
index 4301dfc..4bfe88f 100644
--- a/Modules/_testinternalcapi.c
+++ b/Modules/_testinternalcapi.c
@@ -1000,6 +1000,172 @@ get_co_localskinds(PyObject *self, PyObject *arg)
 }
 
 static PyObject *
+get_code_var_counts(PyObject *self, PyObject *_args, PyObject *_kwargs)
+{
+    PyThreadState *tstate = _PyThreadState_GET();
+    PyObject *codearg;
+    PyObject *globalnames = NULL;
+    PyObject *attrnames = NULL;
+    PyObject *globalsns = NULL;
+    PyObject *builtinsns = NULL;
+    static char *kwlist[] = {"code", "globalnames", "attrnames", "globalsns",
+                             "builtinsns", NULL};
+    if (!PyArg_ParseTupleAndKeywords(_args, _kwargs,
+                    "O|OOO!O!:get_code_var_counts", kwlist,
+                    &codearg, &globalnames, &attrnames,
+                    &PyDict_Type, &globalsns, &PyDict_Type, &builtinsns))
+    {
+        return NULL;
+    }
+    if (PyFunction_Check(codearg)) {
+        if (globalsns == NULL) {
+            globalsns = PyFunction_GET_GLOBALS(codearg);
+        }
+        if (builtinsns == NULL) {
+            builtinsns = PyFunction_GET_BUILTINS(codearg);
+        }
+        codearg = PyFunction_GET_CODE(codearg);
+    }
+    else if (!PyCode_Check(codearg)) {
+        PyErr_SetString(PyExc_TypeError,
+                        "argument must be a code object or a function");
+        return NULL;
+    }
+    PyCodeObject *code = (PyCodeObject *)codearg;
+
+    _PyCode_var_counts_t counts = {0};
+    _PyCode_GetVarCounts(code, &counts);
+    if (_PyCode_SetUnboundVarCounts(
+            tstate, code, &counts, globalnames, attrnames,
+            globalsns, builtinsns) < 0)
+    {
+        return NULL;
+    }
+
+#define SET_COUNT(DICT, STRUCT, NAME) \
+    do { \
+        PyObject *count = PyLong_FromLong(STRUCT.NAME); \
+        int res = PyDict_SetItemString(DICT, #NAME, count); \
+        Py_DECREF(count); \
+        if (res < 0) { \
+            goto error; \
+        } \
+    } while (0)
+
+    PyObject *locals = NULL;
+    PyObject *args = NULL;
+    PyObject *cells = NULL;
+    PyObject *hidden = NULL;
+    PyObject *unbound = NULL;
+    PyObject *globals = NULL;
+    PyObject *countsobj = PyDict_New();
+    if (countsobj == NULL) {
+        return NULL;
+    }
+    SET_COUNT(countsobj, counts, total);
+
+    // locals
+    locals = PyDict_New();
+    if (locals == NULL) {
+        goto error;
+    }
+    if (PyDict_SetItemString(countsobj, "locals", locals) < 0) {
+        goto error;
+    }
+    SET_COUNT(locals, counts.locals, total);
+
+    // locals.args
+    args = PyDict_New();
+    if (args == NULL) {
+        goto error;
+    }
+    if (PyDict_SetItemString(locals, "args", args) < 0) {
+        goto error;
+    }
+    SET_COUNT(args, counts.locals.args, total);
+    SET_COUNT(args, counts.locals.args, numposonly);
+    SET_COUNT(args, counts.locals.args, numposorkw);
+    SET_COUNT(args, counts.locals.args, numkwonly);
+    SET_COUNT(args, counts.locals.args, varargs);
+    SET_COUNT(args, counts.locals.args, varkwargs);
+
+    // locals.numpure
+    SET_COUNT(locals, counts.locals, numpure);
+
+    // locals.cells
+    cells = PyDict_New();
+    if (cells == NULL) {
+        goto error;
+    }
+    if (PyDict_SetItemString(locals, "cells", cells) < 0) {
+        goto error;
+    }
+    SET_COUNT(cells, counts.locals.cells, total);
+    SET_COUNT(cells, counts.locals.cells, numargs);
+    SET_COUNT(cells, counts.locals.cells, numothers);
+
+    // locals.hidden
+    hidden = PyDict_New();
+    if (hidden == NULL) {
+        goto error;
+    }
+    if (PyDict_SetItemString(locals, "hidden", hidden) < 0) {
+        goto error;
+    }
+    SET_COUNT(hidden, counts.locals.hidden, total);
+    SET_COUNT(hidden, counts.locals.hidden, numpure);
+    SET_COUNT(hidden, counts.locals.hidden, numcells);
+
+    // numfree
+    SET_COUNT(countsobj, counts, numfree);
+
+    // unbound
+    unbound = PyDict_New();
+    if (unbound == NULL) {
+        goto error;
+    }
+    if (PyDict_SetItemString(countsobj, "unbound", unbound) < 0) {
+        goto error;
+    }
+    SET_COUNT(unbound, counts.unbound, total);
+    SET_COUNT(unbound, counts.unbound, numattrs);
+    SET_COUNT(unbound, counts.unbound, numunknown);
+
+    // unbound.globals
+    globals = PyDict_New();
+    if (globals == NULL) {
+        goto error;
+    }
+    if (PyDict_SetItemString(unbound, "globals", globals) < 0) {
+        goto error;
+    }
+    SET_COUNT(globals, counts.unbound.globals, total);
+    SET_COUNT(globals, counts.unbound.globals, numglobal);
+    SET_COUNT(globals, counts.unbound.globals, numbuiltin);
+    SET_COUNT(globals, counts.unbound.globals, numunknown);
+
+#undef SET_COUNT
+
+    Py_DECREF(locals);
+    Py_DECREF(args);
+    Py_DECREF(cells);
+    Py_DECREF(hidden);
+    Py_DECREF(unbound);
+    Py_DECREF(globals);
+    return countsobj;
+
+error:
+    Py_DECREF(countsobj);
+    Py_XDECREF(locals);
+    Py_XDECREF(args);
+    Py_XDECREF(cells);
+    Py_XDECREF(hidden);
+    Py_XDECREF(unbound);
+    Py_XDECREF(globals);
+    return NULL;
+}
+
+static PyObject *
 jit_enabled(PyObject *self, PyObject *arg)
 {
     return PyBool_FromLong(_PyInterpreterState_GET()->jit);
@@ -2120,6 +2286,8 @@ static PyMethodDef module_functions[] = {
     {"code_returns_only_none", code_returns_only_none, METH_O, NULL},
     {"get_co_framesize", get_co_framesize, METH_O, NULL},
     {"get_co_localskinds", get_co_localskinds, METH_O, NULL},
+    {"get_code_var_counts", _PyCFunction_CAST(get_code_var_counts),
+     METH_VARARGS | METH_KEYWORDS, NULL},
     {"jit_enabled", jit_enabled,  METH_NOARGS, NULL},
 #ifdef _Py_TIER2
     {"add_executor_dependency", add_executor_dependency, METH_VARARGS, NULL},
diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index bf24a4a..d643eb9 100644
--- a/Objects/codeobject.c
+++ b/Objects/codeobject.c
@@ -1690,6 +1690,241 @@ PyCode_GetFreevars(PyCodeObject *code)
 }
 
 
+static int
+identify_unbound_names(PyThreadState *tstate, PyCodeObject *co,
+                       PyObject *globalnames, PyObject *attrnames,
+                       PyObject *globalsns, PyObject *builtinsns,
+                       struct co_unbound_counts *counts)
+{
+    // This function is inspired by inspect.getclosurevars().
+    // It would be nicer if we had something similar to co_localspluskinds,
+    // but for co_names.
+    assert(globalnames != NULL);
+    assert(PySet_Check(globalnames));
+    assert(PySet_GET_SIZE(globalnames) == 0 || counts != NULL);
+    assert(attrnames != NULL);
+    assert(PySet_Check(attrnames));
+    assert(PySet_GET_SIZE(attrnames) == 0 || counts != NULL);
+    assert(globalsns == NULL || PyDict_Check(globalsns));
+    assert(builtinsns == NULL || PyDict_Check(builtinsns));
+    assert(counts == NULL || counts->total == 0);
+    Py_ssize_t len = Py_SIZE(co);
+    for (int i = 0; i < len; i++) {
+        _Py_CODEUNIT inst = _Py_GetBaseCodeUnit(co, i);
+        if (inst.op.code == LOAD_ATTR) {
+            PyObject *name = PyTuple_GET_ITEM(co->co_names, inst.op.arg>>1);
+            if (counts != NULL) {
+                if (PySet_Contains(attrnames, name)) {
+                    if (_PyErr_Occurred(tstate)) {
+                        return -1;
+                    }
+                    continue;
+                }
+                counts->total += 1;
+                counts->numattrs += 1;
+            }
+            if (PySet_Add(attrnames, name) < 0) {
+                return -1;
+            }
+        }
+        else if (inst.op.code == LOAD_GLOBAL) {
+            PyObject *name = PyTuple_GET_ITEM(co->co_names, inst.op.arg>>1);
+            if (counts != NULL) {
+                if (PySet_Contains(globalnames, name)) {
+                    if (_PyErr_Occurred(tstate)) {
+                        return -1;
+                    }
+                    continue;
+                }
+                counts->total += 1;
+                counts->globals.total += 1;
+                counts->globals.numunknown += 1;
+                if (globalsns != NULL && PyDict_Contains(globalsns, name)) {
+                    if (_PyErr_Occurred(tstate)) {
+                        return -1;
+                    }
+                    counts->globals.numglobal += 1;
+                    counts->globals.numunknown -= 1;
+                }
+                if (builtinsns != NULL && PyDict_Contains(builtinsns, name)) {
+                    if (_PyErr_Occurred(tstate)) {
+                        return -1;
+                    }
+                    counts->globals.numbuiltin += 1;
+                    counts->globals.numunknown -= 1;
+                }
+            }
+            if (PySet_Add(globalnames, name) < 0) {
+                return -1;
+            }
+        }
+    }
+    return 0;
+}
+
+
+void
+_PyCode_GetVarCounts(PyCodeObject *co, _PyCode_var_counts_t *counts)
+{
+    // Count the locals, cells, and free vars.
+    struct co_locals_counts locals = {0};
+    int numfree = 0;
+    PyObject *kinds = co->co_localspluskinds;
+    Py_ssize_t numlocalplusfree = PyBytes_GET_SIZE(kinds);
+    for (int i = 0; i < numlocalplusfree; i++) {
+        _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);
+        if (kind & CO_FAST_FREE) {
+            assert(!(kind & CO_FAST_LOCAL));
+            assert(!(kind & CO_FAST_HIDDEN));
+            assert(!(kind & CO_FAST_ARG));
+            numfree += 1;
+        }
+        else {
+            // Apparently not all non-free vars a CO_FAST_LOCAL.
+            assert(kind);
+            locals.total += 1;
+            if (kind & CO_FAST_ARG) {
+                locals.args.total += 1;
+                if (kind & CO_FAST_ARG_VAR) {
+                    if (kind & CO_FAST_ARG_POS) {
+                        assert(!(kind & CO_FAST_ARG_KW));
+                        assert(!locals.args.varargs);
+                        locals.args.varargs = 1;
+                    }
+                    else {
+                        assert(kind & CO_FAST_ARG_KW);
+                        assert(!locals.args.varkwargs);
+                        locals.args.varkwargs = 1;
+                    }
+                }
+                else if (kind & CO_FAST_ARG_POS) {
+                    if (kind & CO_FAST_ARG_KW) {
+                        locals.args.numposorkw += 1;
+                    }
+                    else {
+                        locals.args.numposonly += 1;
+                    }
+                }
+                else {
+                    assert(kind & CO_FAST_ARG_KW);
+                    locals.args.numkwonly += 1;
+                }
+                if (kind & CO_FAST_CELL) {
+                    locals.cells.total += 1;
+                    locals.cells.numargs += 1;
+                }
+                // Args are never hidden currently.
+                assert(!(kind & CO_FAST_HIDDEN));
+            }
+            else {
+                if (kind & CO_FAST_CELL) {
+                    locals.cells.total += 1;
+                    locals.cells.numothers += 1;
+                    if (kind & CO_FAST_HIDDEN) {
+                        locals.hidden.total += 1;
+                        locals.hidden.numcells += 1;
+                    }
+                }
+                else {
+                    locals.numpure += 1;
+                    if (kind & CO_FAST_HIDDEN) {
+                        locals.hidden.total += 1;
+                        locals.hidden.numpure += 1;
+                    }
+                }
+            }
+        }
+    }
+    assert(locals.args.total == (
+            co->co_argcount + co->co_kwonlyargcount
+            + !!(co->co_flags & CO_VARARGS)
+            + !!(co->co_flags & CO_VARKEYWORDS)));
+    assert(locals.args.numposonly == co->co_posonlyargcount);
+    assert(locals.args.numposonly + locals.args.numposorkw == co->co_argcount);
+    assert(locals.args.numkwonly == co->co_kwonlyargcount);
+    assert(locals.cells.total == co->co_ncellvars);
+    assert(locals.args.total + locals.numpure == co->co_nlocals);
+    assert(locals.total + locals.cells.numargs == co->co_nlocals + co->co_ncellvars);
+    assert(locals.total + numfree == co->co_nlocalsplus);
+    assert(numfree == co->co_nfreevars);
+
+    // Get the unbound counts.
+    assert(PyTuple_GET_SIZE(co->co_names) >= 0);
+    struct co_unbound_counts unbound = {
+        .total = (int)PyTuple_GET_SIZE(co->co_names),
+        // numglobal and numattrs can be set later
+        // with _PyCode_SetUnboundVarCounts().
+        .numunknown = (int)PyTuple_GET_SIZE(co->co_names),
+    };
+
+    // "Return" the result.
+    *counts = (_PyCode_var_counts_t){
+        .total = locals.total + numfree + unbound.total,
+        .locals = locals,
+        .numfree = numfree,
+        .unbound = unbound,
+    };
+}
+
+int
+_PyCode_SetUnboundVarCounts(PyThreadState *tstate,
+                            PyCodeObject *co, _PyCode_var_counts_t *counts,
+                            PyObject *globalnames, PyObject *attrnames,
+                            PyObject *globalsns, PyObject *builtinsns)
+{
+    int res = -1;
+    PyObject *globalnames_owned = NULL;
+    PyObject *attrnames_owned = NULL;
+
+    // Prep the name sets.
+    if (globalnames == NULL) {
+        globalnames_owned = PySet_New(NULL);
+        if (globalnames_owned == NULL) {
+            goto finally;
+        }
+        globalnames = globalnames_owned;
+    }
+    else if (!PySet_Check(globalnames)) {
+        _PyErr_Format(tstate, PyExc_TypeError,
+                     "expected a set for \"globalnames\", got %R", globalnames);
+        goto finally;
+    }
+    if (attrnames == NULL) {
+        attrnames_owned = PySet_New(NULL);
+        if (attrnames_owned == NULL) {
+            goto finally;
+        }
+        attrnames = attrnames_owned;
+    }
+    else if (!PySet_Check(attrnames)) {
+        _PyErr_Format(tstate, PyExc_TypeError,
+                     "expected a set for \"attrnames\", got %R", attrnames);
+        goto finally;
+    }
+
+    // Fill in unbound.globals and unbound.numattrs.
+    struct co_unbound_counts unbound = {0};
+    if (identify_unbound_names(
+            tstate, co, globalnames, attrnames, globalsns, builtinsns,
+            &unbound) < 0)
+    {
+        goto finally;
+    }
+    assert(unbound.numunknown == 0);
+    assert(unbound.total <= counts->unbound.total);
+    assert(counts->unbound.numunknown == counts->unbound.total);
+    unbound.numunknown = counts->unbound.total - unbound.total;
+    unbound.total = counts->unbound.total;
+    counts->unbound = unbound;
+    res = 0;
+
+finally:
+    Py_XDECREF(globalnames_owned);
+    Py_XDECREF(attrnames_owned);
+    return res;
+}
+
+
 /* Here "value" means a non-None value, since a bare return is identical
  * to returning None explicitly.  Likewise a missing return statement
  * at the end of the function is turned into "return None". */