gh-76785: Support Running Some Functions in Subinterpreters (gh-110251)

This specifically refers to `test.support.interpreters.Interpreter.run()`.

3 files changed, 439 insertions, 26 deletions

diff --git a/Lib/test/support/interpreters.py b/Lib/test/support/interpreters.py
index 3b50161..5aba369 100644
--- a/Lib/test/support/interpreters.py
+++ b/Lib/test/support/interpreters.py
@@ -91,12 +91,26 @@ class Interpreter:
         """
         return _interpreters.destroy(self._id)
 
+    # XXX Rename "run" to "exec"?
     def run(self, src_str, /, *, channels=None):
         """Run the given source code in the interpreter.
 
-        This blocks the current Python thread until done.
+        This is essentially the same as calling the builtin "exec"
+        with this interpreter, using the __dict__ of its __main__
+        module as both globals and locals.
+
+        There is no return value.
+
+        If the code raises an unhandled exception then a RunFailedError
+        is raised, which summarizes the unhandled exception.  The actual
+        exception is discarded because objects cannot be shared between
+        interpreters.
+
+        This blocks the current Python thread until done.  During
+        that time, the previous interpreter is allowed to run
+        in other threads.
         """
-        _interpreters.run_string(self._id, src_str, channels)
+        _interpreters.exec(self._id, src_str, channels)
 
 
 def create_channel():
diff --git a/Lib/test/test__xxsubinterpreters.py b/Lib/test/test__xxsubinterpreters.py
index ac2280e..e3c917a 100644
--- a/Lib/test/test__xxsubinterpreters.py
+++ b/Lib/test/test__xxsubinterpreters.py
@@ -925,5 +925,110 @@ class RunStringTests(TestBase):
         self.assertEqual(retcode, 0)
 
 
+class RunFuncTests(TestBase):
+
+    def setUp(self):
+        super().setUp()
+        self.id = interpreters.create()
+
+    def test_success(self):
+        r, w = os.pipe()
+        def script():
+            global w
+            import contextlib
+            with open(w, 'w', encoding="utf-8") as spipe:
+                with contextlib.redirect_stdout(spipe):
+                    print('it worked!', end='')
+        interpreters.run_func(self.id, script, shared=dict(w=w))
+
+        with open(r, encoding="utf-8") as outfile:
+            out = outfile.read()
+
+        self.assertEqual(out, 'it worked!')
+
+    def test_in_thread(self):
+        r, w = os.pipe()
+        def script():
+            global w
+            import contextlib
+            with open(w, 'w', encoding="utf-8") as spipe:
+                with contextlib.redirect_stdout(spipe):
+                    print('it worked!', end='')
+        def f():
+            interpreters.run_func(self.id, script, shared=dict(w=w))
+        t = threading.Thread(target=f)
+        t.start()
+        t.join()
+
+        with open(r, encoding="utf-8") as outfile:
+            out = outfile.read()
+
+        self.assertEqual(out, 'it worked!')
+
+    def test_code_object(self):
+        r, w = os.pipe()
+
+        def script():
+            global w
+            import contextlib
+            with open(w, 'w', encoding="utf-8") as spipe:
+                with contextlib.redirect_stdout(spipe):
+                    print('it worked!', end='')
+        code = script.__code__
+        interpreters.run_func(self.id, code, shared=dict(w=w))
+
+        with open(r, encoding="utf-8") as outfile:
+            out = outfile.read()
+
+        self.assertEqual(out, 'it worked!')
+
+    def test_closure(self):
+        spam = True
+        def script():
+            assert spam
+
+        with self.assertRaises(ValueError):
+            interpreters.run_func(self.id, script)
+
+    # XXX This hasn't been fixed yet.
+    @unittest.expectedFailure
+    def test_return_value(self):
+        def script():
+            return 'spam'
+        with self.assertRaises(ValueError):
+            interpreters.run_func(self.id, script)
+
+    def test_args(self):
+        with self.subTest('args'):
+            def script(a, b=0):
+                assert a == b
+            with self.assertRaises(ValueError):
+                interpreters.run_func(self.id, script)
+
+        with self.subTest('*args'):
+            def script(*args):
+                assert not args
+            with self.assertRaises(ValueError):
+                interpreters.run_func(self.id, script)
+
+        with self.subTest('**kwargs'):
+            def script(**kwargs):
+                assert not kwargs
+            with self.assertRaises(ValueError):
+                interpreters.run_func(self.id, script)
+
+        with self.subTest('kwonly'):
+            def script(*, spam=True):
+                assert spam
+            with self.assertRaises(ValueError):
+                interpreters.run_func(self.id, script)
+
+        with self.subTest('posonly'):
+            def script(spam, /):
+                assert spam
+            with self.assertRaises(ValueError):
+                interpreters.run_func(self.id, script)
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/Modules/_xxsubinterpretersmodule.c b/Modules/_xxsubinterpretersmodule.c
index bca16ac..12c98ea 100644
--- a/Modules/_xxsubinterpretersmodule.c
+++ b/Modules/_xxsubinterpretersmodule.c
@@ -10,6 +10,7 @@
 #include "pycore_pyerrors.h"      // _PyErr_ChainExceptions1()
 #include "pycore_pystate.h"       // _PyInterpreterState_SetRunningMain()
 #include "interpreteridobject.h"
+#include "marshal.h"              // PyMarshal_ReadObjectFromString()
 
 
 #define MODULE_NAME "_xxsubinterpreters"
@@ -366,6 +367,98 @@ _sharedexception_apply(_sharedexception *exc, PyObject *wrapperclass)
 }
 
 
+/* Python code **************************************************************/
+
+static const char *
+check_code_str(PyUnicodeObject *text)
+{
+    assert(text != NULL);
+    if (PyUnicode_GET_LENGTH(text) == 0) {
+        return "too short";
+    }
+
+    // XXX Verify that it parses?
+
+    return NULL;
+}
+
+static const char *
+check_code_object(PyCodeObject *code)
+{
+    assert(code != NULL);
+    if (code->co_argcount > 0
+        || code->co_posonlyargcount > 0
+        || code->co_kwonlyargcount > 0
+        || code->co_flags & (CO_VARARGS | CO_VARKEYWORDS))
+    {
+        return "arguments not supported";
+    }
+    if (code->co_ncellvars > 0) {
+        return "closures not supported";
+    }
+    // We trust that no code objects under co_consts have unbound cell vars.
+
+    if (code->co_executors != NULL
+        || code->_co_instrumentation_version > 0)
+    {
+        return "only basic functions are supported";
+    }
+    if (code->_co_monitoring != NULL) {
+        return "only basic functions are supported";
+    }
+    if (code->co_extra != NULL) {
+        return "only basic functions are supported";
+    }
+
+    return NULL;
+}
+
+#define RUN_TEXT 1
+#define RUN_CODE 2
+
+static const char *
+get_code_str(PyObject *arg, Py_ssize_t *len_p, PyObject **bytes_p, int *flags_p)
+{
+    const char *codestr = NULL;
+    Py_ssize_t len = -1;
+    PyObject *bytes_obj = NULL;
+    int flags = 0;
+
+    if (PyUnicode_Check(arg)) {
+        assert(PyUnicode_CheckExact(arg)
+               && (check_code_str((PyUnicodeObject *)arg) == NULL));
+        codestr = PyUnicode_AsUTF8AndSize(arg, &len);
+        if (codestr == NULL) {
+            return NULL;
+        }
+        if (strlen(codestr) != (size_t)len) {
+            PyErr_SetString(PyExc_ValueError,
+                            "source code string cannot contain null bytes");
+            return NULL;
+        }
+        flags = RUN_TEXT;
+    }
+    else {
+        assert(PyCode_Check(arg)
+               && (check_code_object((PyCodeObject *)arg) == NULL));
+        flags = RUN_CODE;
+
+        // Serialize the code object.
+        bytes_obj = PyMarshal_WriteObjectToString(arg, Py_MARSHAL_VERSION);
+        if (bytes_obj == NULL) {
+            return NULL;
+        }
+        codestr = PyBytes_AS_STRING(bytes_obj);
+        len = PyBytes_GET_SIZE(bytes_obj);
+    }
+
+    *flags_p = flags;
+    *bytes_p = bytes_obj;
+    *len_p = len;
+    return codestr;
+}
+
+
 /* interpreter-specific code ************************************************/
 
 static int
@@ -393,8 +486,9 @@ exceptions_init(PyObject *mod)
 }
 
 static int
-_run_script(PyInterpreterState *interp, const char *codestr,
-            _sharedns *shared, _sharedexception *sharedexc)
+_run_script(PyInterpreterState *interp,
+            const char *codestr, Py_ssize_t codestrlen,
+            _sharedns *shared, _sharedexception *sharedexc, int flags)
 {
     int errcode = ERR_NOT_SET;
 
@@ -428,8 +522,21 @@ _run_script(PyInterpreterState *interp, const char *codestr,
         }
     }
 
-    // Run the string (see PyRun_SimpleStringFlags).
-    PyObject *result = PyRun_StringFlags(codestr, Py_file_input, ns, ns, NULL);
+    // Run the script/code/etc.
+    PyObject *result = NULL;
+    if (flags & RUN_TEXT) {
+        result = PyRun_StringFlags(codestr, Py_file_input, ns, ns, NULL);
+    }
+    else if (flags & RUN_CODE) {
+        PyObject *code = PyMarshal_ReadObjectFromString(codestr, codestrlen);
+        if (code != NULL) {
+            result = PyEval_EvalCode(code, ns, ns);
+            Py_DECREF(code);
+        }
+    }
+    else {
+        Py_UNREACHABLE();
+    }
     Py_DECREF(ns);
     if (result == NULL) {
         goto error;
@@ -465,8 +572,9 @@ error:
 }
 
 static int
-_run_script_in_interpreter(PyObject *mod, PyInterpreterState *interp,
-                           const char *codestr, PyObject *shareables)
+_run_in_interpreter(PyObject *mod, PyInterpreterState *interp,
+                    const char *codestr, Py_ssize_t codestrlen,
+                    PyObject *shareables, int flags)
 {
     module_state *state = get_module_state(mod);
     assert(state != NULL);
@@ -488,7 +596,7 @@ _run_script_in_interpreter(PyObject *mod, PyInterpreterState *interp,
 
     // Run the script.
     _sharedexception exc = (_sharedexception){ .interp = interp };
-    int result = _run_script(interp, codestr, shared, &exc);
+    int result = _run_script(interp, codestr, codestrlen, shared, &exc, flags);
 
     // Switch back.
     if (save_tstate != NULL) {
@@ -695,49 +803,231 @@ PyDoc_STRVAR(get_main_doc,
 Return the ID of main interpreter.");
 
 
+static PyUnicodeObject *
+convert_script_arg(PyObject *arg, const char *fname, const char *displayname,
+                   const char *expected)
+{
+    PyUnicodeObject *str = NULL;
+    if (PyUnicode_CheckExact(arg)) {
+        str = (PyUnicodeObject *)Py_NewRef(arg);
+    }
+    else if (PyUnicode_Check(arg)) {
+        // XXX str = PyUnicode_FromObject(arg);
+        str = (PyUnicodeObject *)Py_NewRef(arg);
+    }
+    else {
+        _PyArg_BadArgument(fname, displayname, expected, arg);
+        return NULL;
+    }
+
+    const char *err = check_code_str(str);
+    if (err != NULL) {
+        Py_DECREF(str);
+        PyErr_Format(PyExc_ValueError,
+                     "%.200s(): bad script text (%s)", fname, err);
+        return NULL;
+    }
+
+    return str;
+}
+
+static PyCodeObject *
+convert_code_arg(PyObject *arg, const char *fname, const char *displayname,
+                 const char *expected)
+{
+    const char *kind = NULL;
+    PyCodeObject *code = NULL;
+    if (PyFunction_Check(arg)) {
+        if (PyFunction_GetClosure(arg) != NULL) {
+            PyErr_Format(PyExc_ValueError,
+                         "%.200s(): closures not supported", fname);
+            return NULL;
+        }
+        code = (PyCodeObject *)PyFunction_GetCode(arg);
+        if (code == NULL) {
+            if (PyErr_Occurred()) {
+                // This chains.
+                PyErr_Format(PyExc_ValueError,
+                             "%.200s(): bad func", fname);
+            }
+            else {
+                PyErr_Format(PyExc_ValueError,
+                             "%.200s(): func.__code__ missing", fname);
+            }
+            return NULL;
+        }
+        Py_INCREF(code);
+        kind = "func";
+    }
+    else if (PyCode_Check(arg)) {
+        code = (PyCodeObject *)Py_NewRef(arg);
+        kind = "code object";
+    }
+    else {
+        _PyArg_BadArgument(fname, displayname, expected, arg);
+        return NULL;
+    }
+
+    const char *err = check_code_object(code);
+    if (err != NULL) {
+        Py_DECREF(code);
+        PyErr_Format(PyExc_ValueError,
+                     "%.200s(): bad %s (%s)", fname, kind, err);
+        return NULL;
+    }
+
+    return code;
+}
+
+static int
+_interp_exec(PyObject *self,
+             PyObject *id_arg, PyObject *code_arg, PyObject *shared_arg)
+{
+    // Look up the interpreter.
+    PyInterpreterState *interp = PyInterpreterID_LookUp(id_arg);
+    if (interp == NULL) {
+        return -1;
+    }
+
+    // Extract code.
+    Py_ssize_t codestrlen = -1;
+    PyObject *bytes_obj = NULL;
+    int flags = 0;
+    const char *codestr = get_code_str(code_arg,
+                                       &codestrlen, &bytes_obj, &flags);
+    if (codestr == NULL) {
+        return -1;
+    }
+
+    // Run the code in the interpreter.
+    int res = _run_in_interpreter(self, interp, codestr, codestrlen,
+                                  shared_arg, flags);
+    Py_XDECREF(bytes_obj);
+    if (res != 0) {
+        return -1;
+    }
+
+    return 0;
+}
+
 static PyObject *
-interp_run_string(PyObject *self, PyObject *args, PyObject *kwds)
+interp_exec(PyObject *self, PyObject *args, PyObject *kwds)
 {
-    static char *kwlist[] = {"id", "script", "shared", NULL};
+    static char *kwlist[] = {"id", "code", "shared", NULL};
     PyObject *id, *code;
     PyObject *shared = NULL;
     if (!PyArg_ParseTupleAndKeywords(args, kwds,
-                                     "OU|O:run_string", kwlist,
+                                     "OO|O:" MODULE_NAME ".exec", kwlist,
                                      &id, &code, &shared)) {
         return NULL;
     }
 
-    // Look up the interpreter.
-    PyInterpreterState *interp = PyInterpreterID_LookUp(id);
-    if (interp == NULL) {
+    const char *expected = "a string, a function, or a code object";
+    if (PyUnicode_Check(code)) {
+         code = (PyObject *)convert_script_arg(code, MODULE_NAME ".exec",
+                                               "argument 2", expected);
+    }
+    else {
+         code = (PyObject *)convert_code_arg(code, MODULE_NAME ".exec",
+                                             "argument 2", expected);
+    }
+    if (code == NULL) {
         return NULL;
     }
 
-    // Extract code.
-    Py_ssize_t size;
-    const char *codestr = PyUnicode_AsUTF8AndSize(code, &size);
-    if (codestr == NULL) {
+    int res = _interp_exec(self, id, code, shared);
+    Py_DECREF(code);
+    if (res < 0) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+PyDoc_STRVAR(exec_doc,
+"exec(id, code, shared=None)\n\
+\n\
+Execute the provided code in the identified interpreter.\n\
+This is equivalent to running the builtin exec() under the target\n\
+interpreter, using the __dict__ of its __main__ module as both\n\
+globals and locals.\n\
+\n\
+\"code\" may be a string containing the text of a Python script.\n\
+\n\
+Functions (and code objects) are also supported, with some restrictions.\n\
+The code/function must not take any arguments or be a closure\n\
+(i.e. have cell vars).  Methods and other callables are not supported.\n\
+\n\
+If a function is provided, its code object is used and all its state\n\
+is ignored, including its __globals__ dict.");
+
+static PyObject *
+interp_run_string(PyObject *self, PyObject *args, PyObject *kwds)
+{
+    static char *kwlist[] = {"id", "script", "shared", NULL};
+    PyObject *id, *script;
+    PyObject *shared = NULL;
+    if (!PyArg_ParseTupleAndKeywords(args, kwds,
+                                     "OU|O:" MODULE_NAME ".run_string", kwlist,
+                                     &id, &script, &shared)) {
         return NULL;
     }
-    if (strlen(codestr) != (size_t)size) {
-        PyErr_SetString(PyExc_ValueError,
-                        "source code string cannot contain null bytes");
+
+    script = (PyObject *)convert_script_arg(script, MODULE_NAME ".exec",
+                                            "argument 2", "a string");
+    if (script == NULL) {
         return NULL;
     }
 
-    // Run the code in the interpreter.
-    if (_run_script_in_interpreter(self, interp, codestr, shared) != 0) {
+    int res = _interp_exec(self, id, (PyObject *)script, shared);
+    Py_DECREF(script);
+    if (res < 0) {
         return NULL;
     }
     Py_RETURN_NONE;
 }
 
 PyDoc_STRVAR(run_string_doc,
-"run_string(id, script, shared)\n\
+"run_string(id, script, shared=None)\n\
 \n\
 Execute the provided string in the identified interpreter.\n\
 \n\
-See PyRun_SimpleStrings.");
+(See " MODULE_NAME ".exec().");
+
+static PyObject *
+interp_run_func(PyObject *self, PyObject *args, PyObject *kwds)
+{
+    static char *kwlist[] = {"id", "func", "shared", NULL};
+    PyObject *id, *func;
+    PyObject *shared = NULL;
+    if (!PyArg_ParseTupleAndKeywords(args, kwds,
+                                     "OO|O:" MODULE_NAME ".run_func", kwlist,
+                                     &id, &func, &shared)) {
+        return NULL;
+    }
+
+    PyCodeObject *code = convert_code_arg(func, MODULE_NAME ".exec",
+                                          "argument 2",
+                                          "a function or a code object");
+    if (code == NULL) {
+        return NULL;
+    }
+
+    int res = _interp_exec(self, id, (PyObject *)code, shared);
+    Py_DECREF(code);
+    if (res < 0) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+PyDoc_STRVAR(run_func_doc,
+"run_func(id, func, shared=None)\n\
+\n\
+Execute the body of the provided function in the identified interpreter.\n\
+Code objects are also supported.  In both cases, closures and args\n\
+are not supported.  Methods and other callables are not supported either.\n\
+\n\
+(See " MODULE_NAME ".exec().");
 
 
 static PyObject *
@@ -804,8 +1094,12 @@ static PyMethodDef module_functions[] = {
 
     {"is_running",                _PyCFunction_CAST(interp_is_running),
      METH_VARARGS | METH_KEYWORDS, is_running_doc},
+    {"exec",                      _PyCFunction_CAST(interp_exec),
+     METH_VARARGS | METH_KEYWORDS, exec_doc},
     {"run_string",                _PyCFunction_CAST(interp_run_string),
      METH_VARARGS | METH_KEYWORDS, run_string_doc},
+    {"run_func",                  _PyCFunction_CAST(interp_run_func),
+     METH_VARARGS | METH_KEYWORDS, run_func_doc},
 
     {"is_shareable",              _PyCFunction_CAST(object_is_shareable),
      METH_VARARGS | METH_KEYWORDS, is_shareable_doc},