gh-132775: Add _PyCode_ReturnsOnlyNone() (gh-132981)

The function indicates whether or not the function has a return statement.

This is used by a later change related treating some functions like scripts.

Author: ericsnowcurrently

Include/internal/pycore_code.h

@@ -561,6 +561,10 @@ extern void _Py_ClearTLBCIndex(_PyThreadStateImpl *tstate);
 extern int _Py_ClearUnusedTLBC(PyInterpreterState *interp);
 #endif
 
+
+PyAPI_FUNC(int) _PyCode_ReturnsOnlyNone(PyCodeObject *);
+
+
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_opcode_utils.h

@@ -54,6 +54,9 @@ extern "C" {
          (opcode) == RAISE_VARARGS || \
          (opcode) == RERAISE)
 
+#define IS_RETURN_OPCODE(opcode) \
+        (opcode == RETURN_VALUE)
+
 
 /* Flags used in the oparg for MAKE_FUNCTION */
 #define MAKE_FUNCTION_DEFAULTS    0x01

Lib/test/test_code.py

@@ -216,6 +216,10 @@
 from test.support.bytecode_helper import instructions_with_positions
 from opcode import opmap, opname
 from _testcapi import code_offset_to_line
+try:
+    import _testinternalcapi
+except ModuleNotFoundError:
+    _testinternalcapi = None
 
 COPY_FREE_VARS = opmap['COPY_FREE_VARS']
 
@@ -425,6 +429,61 @@ def func():
         with self.assertWarns(DeprecationWarning):
             func.__code__.co_lnotab
 
+    @unittest.skipIf(_testinternalcapi is None, '_testinternalcapi is missing')
+    def test_returns_only_none(self):
+        value = True
+
+        def spam1():
+            pass
+        def spam2():
+            return
+        def spam3():
+            return None
+        def spam4():
+            if not value:
+                return
+            ...
+        def spam5():
+            if not value:
+                return None
+            ...
+        lambda1 = (lambda: None)
+        for func in [
+            spam1,
+            spam2,
+            spam3,
+            spam4,
+            spam5,
+            lambda1,
+        ]:
+            with self.subTest(func):
+                res = _testinternalcapi.code_returns_only_none(func.__code__)
+                self.assertTrue(res)
+
+        def spam6():
+            return True
+        def spam7():
+            return value
+        def spam8():
+            if value:
+                return None
+            return True
+        def spam9():
+            if value:
+                return True
+            return None
+        lambda2 = (lambda: True)
+        for func in [
+            spam6,
+            spam7,
+            spam8,
+            spam9,
+            lambda2,
+        ]:
+            with self.subTest(func):
+                res = _testinternalcapi.code_returns_only_none(func.__code__)
+                self.assertFalse(res)
+
     def test_invalid_bytecode(self):
         def foo():
             pass

Modules/_testinternalcapi.c

@@ -945,6 +945,18 @@ iframe_getlasti(PyObject *self, PyObject *frame)
     return PyLong_FromLong(PyUnstable_InterpreterFrame_GetLasti(f));
 }
 
+static PyObject *
+code_returns_only_none(PyObject *self, PyObject *arg)
+{
+    if (!PyCode_Check(arg)) {
+        PyErr_SetString(PyExc_TypeError, "argument must be a code object");
+        return NULL;
+    }
+    PyCodeObject *code = (PyCodeObject *)arg;
+    int res = _PyCode_ReturnsOnlyNone(code);
+    return PyBool_FromLong(res);
+}
+
 static PyObject *
 get_co_framesize(PyObject *self, PyObject *arg)
 {
@@ -2074,6 +2086,7 @@ static PyMethodDef module_functions[] = {
     {"iframe_getcode", iframe_getcode, METH_O, NULL},
     {"iframe_getline", iframe_getline, METH_O, NULL},
     {"iframe_getlasti", iframe_getlasti, METH_O, NULL},
+    {"code_returns_only_none", code_returns_only_none, METH_O, NULL},
     {"get_co_framesize", get_co_framesize, METH_O, NULL},
     {"jit_enabled", jit_enabled,  METH_NOARGS, NULL},
 #ifdef _Py_TIER2

Objects/codeobject.c

@@ -1689,6 +1689,49 @@ PyCode_GetFreevars(PyCodeObject *code)
     return _PyCode_GetFreevars(code);
 }
 
+
+/* 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". */
+int
+_PyCode_ReturnsOnlyNone(PyCodeObject *co)
+{
+    // Look up None in co_consts.
+    Py_ssize_t nconsts = PyTuple_Size(co->co_consts);
+    int none_index = 0;
+    for (; none_index < nconsts; none_index++) {
+        if (PyTuple_GET_ITEM(co->co_consts, none_index) == Py_None) {
+            break;
+        }
+    }
+    if (none_index == nconsts) {
+        // None wasn't there, which means there was no implicit return,
+        // "return", or "return None".  That means there must be
+        // an explicit return (non-None).
+        return 0;
+    }
+
+    // Walk the bytecode, looking for RETURN_VALUE.
+    Py_ssize_t len = Py_SIZE(co);
+    for (int i = 0; i < len; i++) {
+        _Py_CODEUNIT inst = _Py_GetBaseCodeUnit(co, i);
+        if (IS_RETURN_OPCODE(inst.op.code)) {
+            assert(i != 0);
+            // Ignore it if it returns None.
+            _Py_CODEUNIT prev = _Py_GetBaseCodeUnit(co, i-1);
+            if (prev.op.code == LOAD_CONST) {
+                // We don't worry about EXTENDED_ARG for now.
+                if (prev.op.arg == none_index) {
+                    continue;
+                }
+            }
+            return 0;
+        }
+    }
+    return 1;
+}
+
+
 #ifdef _Py_TIER2
 
 static void

Python/flowgraph.c

@@ -295,7 +295,7 @@ dump_instr(cfg_instr *i)
 static inline int
 basicblock_returns(const basicblock *b) {
     cfg_instr *last = basicblock_last_instr(b);
-    return last && last->i_opcode == RETURN_VALUE;
+    return last && IS_RETURN_OPCODE(last->i_opcode);
 }
 
 static void