aboutsummaryrefslogtreecommitdiffstatshomepage
path: root/Objects/codeobject.c
diff options
context:
space:
mode:
authorEric Snow <ericsnowcurrently@gmail.com>2025-04-30 12:19:20 -0600
committerGitHub <noreply@github.com>2025-04-30 18:19:20 +0000
commit94b4fcd806e7b692955173d309ea3b70a193ad96 (patch)
treee1ac874a97b3fca6f0ff763b1dd5eaaf9c4fca47 /Objects/codeobject.c
parent26c0248b54b6b2a5df51dd3da8c0ebb1b2958bc4 (diff)
downloadcpython-94b4fcd806e7b692955173d309ea3b70a193ad96.tar.gz
cpython-94b4fcd806e7b692955173d309ea3b70a193ad96.zip
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.
Diffstat (limited to 'Objects/codeobject.c')
-rw-r--r--Objects/codeobject.c235
1 files changed, 235 insertions, 0 deletions
diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index bf24a4af445..d643eb9fd61 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". */
generated by cgit v1.2.3 (git 2.39.1) at 2025-06-05 12:04:40 +0000