Optionally identify globals/builtins counts.

Author: ericsnowcurrently

Include/internal/pycore_code.h

@@ -592,7 +592,12 @@ typedef struct {
     int numfree;  // nonlocal
     struct co_unbound_counts {
         int total;
-        int numglobal;
+        struct {
+            int total;
+            int numglobal;
+            int numbuiltin;
+            int numunknown;
+        } globals;
         int numattrs;
         int numunknown;
     } unbound;
@@ -604,8 +609,10 @@ PyAPI_FUNC(void) _PyCode_GetVarCounts(
 PyAPI_FUNC(int) _PyCode_SetUnboundVarCounts(
         PyCodeObject *,
         _PyCode_var_counts_t *,
-        PyObject *global,
-        PyObject *attrs);
+        PyObject *globalnames,
+        PyObject *attrnames,
+        PyObject *globalsns,
+        PyObject *builtinsns);
 
 PyAPI_FUNC(int) _PyCode_ReturnsOnlyNone(PyCodeObject *);
 

Lib/test/test_code.py

@@ -796,7 +796,36 @@ def new_var_counts(*,
                            ):
             nargvars = posonly + posorkw + kwonly + varargs + varkwargs
             nlocals = nargvars + purelocals + othercells
-            unbound = globalvars + attrs + unknown
+            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': {
@@ -824,7 +853,7 @@ def new_var_counts(*,
                 'numfree': freevars,
                 'unbound': {
                     'total': unbound,
-                    'numglobal': globalvars,
+                    'globals': globalvars,
                     'numattrs': attrs,
                     'numunknown': unknown,
                 },
@@ -929,6 +958,55 @@ def new_var_counts(*,
                 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])

Modules/_testinternalcapi.c

@@ -1002,19 +1002,41 @@ get_co_localskinds(PyObject *self, PyObject *arg)
 static PyObject *
 get_code_var_counts(PyObject *self, PyObject *_args, PyObject *_kwargs)
 {
-    const char *codearg;
-    static char *kwlist[] = {"code", NULL};
+    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!:get_code_var_counts", kwlist,
-                    &PyCode_Type, &codearg))
+                    "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(code, &counts, NULL, NULL) < 0) {
+    if (_PyCode_SetUnboundVarCounts(
+            code, &counts, globalnames, attrnames, globalsns, builtinsns) < 0)
+    {
         return NULL;
     }
 
@@ -1033,6 +1055,7 @@ get_code_var_counts(PyObject *self, PyObject *_args, PyObject *_kwargs)
     PyObject *cells = NULL;
     PyObject *hidden = NULL;
     PyObject *unbound = NULL;
+    PyObject *globals = NULL;
     PyObject *countsobj = PyDict_New();
     if (countsobj == NULL) {
         return NULL;
@@ -1103,17 +1126,30 @@ get_code_var_counts(PyObject *self, PyObject *_args, PyObject *_kwargs)
         goto error;
     }
     SET_COUNT(unbound, counts.unbound, total);
-    SET_COUNT(unbound, counts.unbound, numglobal);
     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:
@@ -1123,6 +1159,7 @@ get_code_var_counts(PyObject *self, PyObject *_args, PyObject *_kwargs)
     Py_XDECREF(cells);
     Py_XDECREF(hidden);
     Py_XDECREF(unbound);
+    Py_XDECREF(globals);
     return NULL;
 }
 

Objects/codeobject.c

@@ -1692,26 +1692,29 @@ PyCode_GetFreevars(PyCodeObject *code)
 
 static int
 identify_unbound_names(PyCodeObject *co,
-                       PyObject *global, PyObject *attrs,
+                       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(global != NULL);
-    assert(PySet_Check(global));
-    assert(PySet_GET_SIZE(global) == 0 || counts != NULL);
-    assert(attrs != NULL);
-    assert(PySet_Check(attrs));
-    assert(PySet_GET_SIZE(attrs) == 0 || counts != NULL);
+    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(attrs, name)) {
+                if (PySet_Contains(attrnames, name)) {
                     if (PyErr_Occurred()) {
                         return -1;
                     }
@@ -1720,23 +1723,38 @@ identify_unbound_names(PyCodeObject *co,
                 counts->total += 1;
                 counts->numattrs += 1;
             }
-            if (PySet_Add(attrs, name) < 0) {
+            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(global, name)) {
+                if (PySet_Contains(globalnames, name)) {
                     if (PyErr_Occurred()) {
                         return -1;
                     }
                     continue;
                 }
                 counts->total += 1;
-                counts->numglobal += 1;
+                counts->globals.total += 1;
+                counts->globals.numunknown += 1;
+                if (globalsns != NULL && PyDict_Contains(globalsns, name)) {
+                    if (PyErr_Occurred()) {
+                        return -1;
+                    }
+                    counts->globals.numglobal += 1;
+                    counts->globals.numunknown -= 1;
+                }
+                if (builtinsns != NULL && PyDict_Contains(builtinsns, name)) {
+                    if (PyErr_Occurred()) {
+                        return -1;
+                    }
+                    counts->globals.numbuiltin += 1;
+                    counts->globals.numunknown -= 1;
+                }
             }
-            if (PySet_Add(global, name) < 0) {
+            if (PySet_Add(globalnames, name) < 0) {
                 return -1;
             }
         }
@@ -1849,46 +1867,41 @@ _PyCode_GetVarCounts(PyCodeObject *co, _PyCode_var_counts_t *counts)
 
 int
 _PyCode_SetUnboundVarCounts(PyCodeObject *co, _PyCode_var_counts_t *counts,
-                            PyObject *globalarg, PyObject *attrsarg)
+                            PyObject *globalnames, PyObject *attrnames,
+                            PyObject *globalsns, PyObject *builtinsns)
 {
     int res = -1;
-    PyObject *global = NULL;
-    PyObject *attrs = NULL;
-    PyObject *global_owned = NULL;
-    PyObject *attrs_owned = NULL;
-    if (globalarg != NULL) {
-        if (!PySet_Check(globalarg)) {
-            PyErr_Format(PyExc_TypeError,
-                         "expected a set for \"global\", got %R", global);
+    PyObject *globalnames_owned = NULL;
+    PyObject *attrnames_owned = NULL;
+    if (globalnames == NULL) {
+        globalnames_owned = PySet_New(NULL);
+        if (globalnames_owned == NULL) {
             goto finally;
         }
-        global = globalarg;
+        globalnames = globalnames_owned;
     }
-    else {
-        global_owned = PySet_New(NULL);
-        if (global_owned == NULL) {
-            goto finally;
-        }
-        global = global_owned;
+    else if (!PySet_Check(globalnames)) {
+        PyErr_Format(PyExc_TypeError,
+                     "expected a set for \"globalnames\", got %R", globalnames);
+        goto finally;
     }
-    if (attrsarg != NULL) {
-        if (!PySet_Check(attrsarg)) {
-            PyErr_Format(PyExc_TypeError,
-                         "expected a set for \"attrs\", got %R", attrs);
+    if (attrnames == NULL) {
+        attrnames_owned = PySet_New(NULL);
+        if (attrnames_owned == NULL) {
             goto finally;
         }
-        attrs = attrsarg;
+        attrnames = attrnames_owned;
     }
-    else {
-        attrs_owned = PySet_New(NULL);
-        if (attrs_owned == NULL) {
-            goto finally;
-        }
-        attrs = attrs_owned;
+    else if (!PySet_Check(attrnames)) {
+        PyErr_Format(PyExc_TypeError,
+                     "expected a set for \"attrnames\", got %R", attrnames);
+        goto finally;
     }
 
     struct co_unbound_counts unbound = {0};
-    if (identify_unbound_names(co, global, attrs, &unbound) < 0) {
+    if (identify_unbound_names(
+            co, globalnames, attrnames, globalsns, builtinsns, &unbound) < 0)
+    {
         goto finally;
     }
     assert(unbound.numunknown == 0);
@@ -1900,8 +1913,8 @@ _PyCode_SetUnboundVarCounts(PyCodeObject *co, _PyCode_var_counts_t *counts,
     res = 0;
 
 finally:
-    Py_XDECREF(global_owned);
-    Py_XDECREF(attrs_owned);
+    Py_XDECREF(globalnames_owned);
+    Py_XDECREF(attrnames_owned);
     return res;
 }