Add _PyCode_GetScriptXIData().

ericsnowcurrently · May 8, 2025 · Apr 24, 2025 · May 2, 2025 · May 8, 2025 · May 8, 2025
commit 49c4abbf201729a010899d2d99e62014a5b9ec13
@@ -191,6 +191,10 @@ PyAPI_FUNC(int) _PyCode_GetXIData(
         PyThreadState *,
         PyObject *,
         _PyXIData_t *);
+PyAPI_FUNC(int) _PyCode_GetScriptXIData(
+        PyThreadState *,
+        PyObject *,
+        _PyXIData_t *);
 
 
 /* using cross-interpreter data */

diff --git a/Lib/test/_code_definitions.py b/Lib/test/_code_definitions.py
@@ -1,4 +1,32 @@
 
+def simple_script():
+    assert True
+
+
+def complex_script():
+    obj = 'a string'
+    pickle = __import__('pickle')
+    def spam_minimal():
+        pass
+    spam_minimal()
+    data = pickle.dumps(obj)
+    res = pickle.loads(data)
+    assert res == obj, (res, obj)
+
+
+def script_with_globals():
+    obj1, obj2 = spam(42)
+    assert obj1 == 42
+    assert obj2 is None
+
+
+def script_with_explicit_empty_return():
+    return None
+
+
+def script_with_return():
+    return True
+
 
 def spam_minimal():
     # no arg defaults or kwarg defaults
@@ -141,6 +169,11 @@ def ham_C_closure(z):
 
 TOP_FUNCTIONS = [
     # shallow
+    simple_script,
+    complex_script,
+    script_with_globals,
+    script_with_explicit_empty_return,
+    script_with_return,
     spam_minimal,
     spam_with_builtins,
     spam_with_globals_and_builtins,
@@ -179,6 +212,10 @@ def ham_C_closure(z):
 ]
 
 STATELESS_FUNCTIONS = [
+    simple_script,
+    complex_script,
+    script_with_explicit_empty_return,
+    script_with_return,
     spam,
     spam_minimal,
     spam_with_builtins,
@@ -200,10 +237,21 @@ def ham_C_closure(z):
 ]
 STATELESS_CODE = [
     *STATELESS_FUNCTIONS,
+    script_with_globals,
     spam_with_globals_and_builtins,
     spam_full,
 ]
 
+SCRIPT_FUNCTIONS = [
+    simple_script,
+    complex_script,
+    script_with_explicit_empty_return,
+    spam_minimal,
+    spam_with_builtins,
+    spam_with_inner_not_closure,
+    spam_with_inner_closure,
+]
+
 
 # generators
 

diff --git a/Lib/test/test_code.py b/Lib/test/test_code.py
@@ -673,6 +673,20 @@ def test_local_kinds(self):
         VARKWARGS = CO_FAST_LOCAL | CO_FAST_ARG_VAR | CO_FAST_ARG_KW
 
         funcs = {
+            defs.simple_script: {},
+            defs.complex_script: {
+                'obj': CO_FAST_LOCAL,
+                'pickle': CO_FAST_LOCAL,
+                'spam_minimal': CO_FAST_LOCAL,
+                'data': CO_FAST_LOCAL,
+                'res': CO_FAST_LOCAL,
+            },
+            defs.script_with_globals: {
+                'obj1': CO_FAST_LOCAL,
+                'obj2': CO_FAST_LOCAL,
+            },
+            defs.script_with_explicit_empty_return: {},
+            defs.script_with_return: {},
             defs.spam_minimal: {},
             defs.spam_with_builtins: {
                 'x': CO_FAST_LOCAL,
@@ -898,6 +912,19 @@ def new_var_counts(*,
             }
 
         funcs = {
+            defs.simple_script: new_var_counts(),
+            defs.complex_script: new_var_counts(
+                purelocals=5,
+                globalvars=1,
+                attrs=2,
+            ),
+            defs.script_with_globals: new_var_counts(
+                purelocals=2,
+                globalvars=1,
+            ),
+            defs.script_with_explicit_empty_return: new_var_counts(),
+            defs.script_with_return: new_var_counts(),
+            defs.spam_minimal: new_var_counts(),
             defs.spam_minimal: new_var_counts(),
             defs.spam_with_builtins: new_var_counts(
                 purelocals=4,

@@ -758,6 +758,97 @@ def test_other_objects(self):
         ])
 
 
+class ShareableScriptTests(_GetXIDataTests):
+
+    MODE = 'script'
+
+    VALID_SCRIPTS = [
+        '',
+        'spam',
+        '# a comment',
+        'print("spam")',
+        'raise Exception("spam")',
+        """if True:
+            do_something()
+            """,
+        """if True:
+            def spam(x):
+                return x
+            class Spam:
+                def eggs(self):
+                    return 42
+            x = Spam().eggs()
+            raise ValueError(spam(x))
+            """,
+    ]
+    INVALID_SCRIPTS = [
+        '    pass',
+        '----',
+        """if True:
+            # do something
+            """,
+    ]
+
+    def test_valid_str(self):
+        self.assert_roundtrip_not_equal([
+            *self.VALID_SCRIPTS,
+        ], expecttype=types.CodeType)
+
+    def test_invalid_str(self):
+        self.assert_not_shareable([
+            *self.INVALID_SCRIPTS,
+        ])
+
+    def test_valid_bytes(self):
+        self.assert_roundtrip_not_equal([
+            *(s.encode('utf8') for s in self.VALID_SCRIPTS),
+        ], expecttype=types.CodeType)
+
+    def test_invalid_bytes(self):
+        self.assert_not_shareable([
+            *(s.encode('utf8') for s in self.INVALID_SCRIPTS),
+        ])
+
+    def test_script_code(self):
+        self.assert_roundtrip_equal_not_identical([
+            *(f.__code__ for f in defs.SCRIPT_FUNCTIONS),
+            defs.script_with_globals.__code__,
+        ])
+
+    def test_other_code(self):
+        self.assert_not_shareable([
+            *(f.__code__ for f in defs.FUNCTIONS
+              if f not in defs.SCRIPT_FUNCTIONS and
+                f is not defs.script_with_globals),
+            *(f.__code__ for f in defs.FUNCTION_LIKE),
+        ])
+
+    def test_script_function(self):
+        self.assert_roundtrip_not_equal([
+            *defs.SCRIPT_FUNCTIONS,
+        ], expecttype=types.CodeType)
+
+    def test_other_function(self):
+        self.assert_not_shareable([
+            *(f for f in defs.FUNCTIONS
+              if f not in defs.SCRIPT_FUNCTIONS),
+            *defs.FUNCTION_LIKE,
+        ])
+
+    def test_other_objects(self):
+        self.assert_not_shareable([
+            None,
+            True,
+            False,
+            Ellipsis,
+            NotImplemented,
+            (),
+            [],
+            {},
+            object(),
+        ])
+
+
 class ShareableTypeTests(_GetXIDataTests):
 
     MODE = 'xidata'

diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
@@ -1989,6 +1989,11 @@ get_crossinterp_data(PyObject *self, PyObject *args, PyObject *kwargs)
             goto error;
         }
     }
+    else if (strcmp(mode, "script") == 0) {
+        if (_PyCode_GetScriptXIData(tstate, obj, xidata) != 0) {
+            goto error;
+        }
+    }
     else {
         PyErr_Format(PyExc_ValueError, "unsupported mode %R", modeobj);
         goto error;

@@ -6,8 +6,10 @@
 #include "osdefs.h"               // MAXPATHLEN
 #include "pycore_ceval.h"         // _Py_simple_func
 #include "pycore_crossinterp.h"   // _PyXIData_t
+#include "pycore_function.h"      // _PyFunction_VerifyStateless()
 #include "pycore_initconfig.h"    // _PyStatus_OK()
 #include "pycore_namespace.h"     // _PyNamespace_New()
+#include "pycore_pythonrun.h"     // _Py_SourceAsString()
 #include "pycore_typeobject.h"    // _PyStaticType_InitBuiltin()
 
 
@@ -784,6 +786,103 @@ _PyMarshal_GetXIData(PyThreadState *tstate, PyObject *obj, _PyXIData_t *xidata)
 }
 
 
+/* script wrapper */
+
+static int
+verify_script(PyThreadState *tstate, PyCodeObject *co)
+{
+    assert(_PyCode_VerifyStateless(
+                tstate, co, NULL, NULL, _PyEval_GetBuiltins(tstate)) == 0);
+    if (co->co_argcount > 0
+        || co->co_posonlyargcount > 0
+        || co->co_kwonlyargcount > 0
+        || co->co_flags & (CO_VARARGS | CO_VARKEYWORDS))
+    {
+        _PyErr_SetString(tstate, PyExc_ValueError,
+                         "code with args not supported");
+        return -1;
+    }
+    if (!_PyCode_ReturnsOnlyNone(co)) {
+        _PyErr_SetString(tstate, PyExc_ValueError,
+                         "code that returns a value is not a script");
+        return -1;
+    }
+    return 0;
+}
+
+int
+_PyCode_GetScriptXIData(PyThreadState *tstate,
+                        PyObject *obj, _PyXIData_t *xidata)
+{
+    // Get the corresponding code object.
+    PyObject *code = NULL;
+    if (PyCode_Check(obj)) {
+        code = obj;
+        Py_INCREF(code);
+        PyObject *builtins = _PyEval_GetBuiltins(tstate);
+        assert(builtins != NULL);
+        if (_PyCode_VerifyStateless(
+                    tstate, (PyCodeObject *)code, NULL, NULL, builtins) < 0)
+        {
+            goto error;
+        }
+        if (verify_script(tstate, (PyCodeObject *)code) < 0) {
+            goto error;
+        }
+    }
+    else if (PyFunction_Check(obj)) {
+        code = PyFunction_GET_CODE(obj);
+        assert(code != NULL);
+        Py_INCREF(code);
+        if (_PyFunction_VerifyStateless(tstate, obj) < 0) {
+            goto error;
+        }
+        if (verify_script(tstate, (PyCodeObject *)code) < 0) {
+            goto error;
+        }
+    }
+    else {
+        const char *filename = "<script>";
+        PyCompilerFlags cf = _PyCompilerFlags_INIT;
+        cf.cf_flags = PyCF_SOURCE_IS_UTF8;
+        PyObject *ref = NULL;
+        const char *script =
+                _Py_SourceAsString(obj, "???", "???", &cf, &ref);
+        if (script == NULL) {
+            if (!PyBytes_Check(obj) && !PyUnicode_Check(obj)
+                && !PyByteArray_Check(obj) && !PyObject_CheckBuffer(obj))
+            {
+                // We discard the raised exception.
+                _PyErr_Format(tstate, PyExc_TypeError, "unsupported script %R", obj);
+            }
+            goto error;
+        }
+        code = Py_CompileStringExFlags(script, filename, Py_file_input, &cf, 0);
+        Py_XDECREF(ref);
+        if (code == NULL) {
+            goto error;
+        }
+    }
+
+    // Convert the code object.
+    int res = _PyCode_GetXIData(tstate, code, xidata);
+    Py_DECREF(code);
+    if (res < 0) {
+        return -1;
+    }
+    return 0;
+
+error:
+    Py_XDECREF(code);
+    PyObject *cause = _PyErr_GetRaisedException(tstate);
+    assert(cause != NULL);
+    _set_xid_lookup_failure(
+                tstate, NULL, "object not a valid script", cause);
+    Py_DECREF(cause);
+    return -1;
+}
+
+
 /* using cross-interpreter data */
 
 PyObject *