gh-132775: Add _PyCode_GetXIData() (gh-133475)

Author: ericsnowcurrently

Include/internal/pycore_crossinterp.h

@@ -185,6 +185,13 @@ PyAPI_FUNC(int) _PyMarshal_GetXIData(
         PyObject *,
         _PyXIData_t *);
 
+// _PyObject_GetXIData() for code objects
+PyAPI_FUNC(PyObject *) _PyCode_FromXIData(_PyXIData_t *);
+PyAPI_FUNC(int) _PyCode_GetXIData(
+        PyThreadState *,
+        PyObject *,
+        _PyXIData_t *);
+
 
 /* using cross-interpreter data */
 

Lib/test/_code_definitions.py

@@ -29,6 +29,12 @@ def spam_with_globals_and_builtins():
     print(res)
 
 
+def spam_args_attrs_and_builtins(a, b, /, c, d, *args, e, f, **kwargs):
+    if args.__len__() > 2:
+        return None
+    return a, b, c, d, e, f, args, kwargs
+
+
 def spam_returns_arg(x):
     return x
 
@@ -46,6 +52,10 @@ def eggs():
     eggs()
 
 
+def spam_annotated(a: int, b: str, c: object) -> tuple:
+    return a, b, c
+
+
 def spam_full(a, b, /, c, d:int=1, *args, e, f:object=None, **kwargs) -> tuple:
     # arg defaults, kwarg defaults
     # annotations
@@ -134,9 +144,11 @@ def ham_C_closure(z):
     spam_minimal,
     spam_with_builtins,
     spam_with_globals_and_builtins,
+    spam_args_attrs_and_builtins,
     spam_returns_arg,
     spam_with_inner_not_closure,
     spam_with_inner_closure,
+    spam_annotated,
     spam_full,
     spam,
     # outer func
@@ -170,7 +182,9 @@ def ham_C_closure(z):
     spam,
     spam_minimal,
     spam_with_builtins,
+    spam_args_attrs_and_builtins,
     spam_returns_arg,
+    spam_annotated,
     spam_with_inner_not_closure,
     spam_with_inner_closure,
     spam_N,

Lib/test/test_code.py

@@ -687,6 +687,16 @@ def test_local_kinds(self):
                 'checks': CO_FAST_LOCAL,
                 'res': CO_FAST_LOCAL,
             },
+            defs.spam_args_attrs_and_builtins: {
+                'a': POSONLY,
+                'b': POSONLY,
+                'c': POSORKW,
+                'd': POSORKW,
+                'e': KWONLY,
+                'f': KWONLY,
+                'args': VARARGS,
+                'kwargs': VARKWARGS,
+            },
             defs.spam_returns_arg: {
                 'x': POSORKW,
             },
@@ -697,6 +707,11 @@ def test_local_kinds(self):
                 'x': CO_FAST_CELL,
                 'eggs': CO_FAST_LOCAL,
             },
+            defs.spam_annotated: {
+                'a': POSORKW,
+                'b': POSORKW,
+                'c': POSORKW,
+            },
             defs.spam_full: {
                 'a': POSONLY,
                 'b': POSONLY,
@@ -892,6 +907,14 @@ def new_var_counts(*,
                 purelocals=5,
                 globalvars=6,
             ),
+            defs.spam_args_attrs_and_builtins: new_var_counts(
+                posonly=2,
+                posorkw=2,
+                kwonly=2,
+                varargs=1,
+                varkwargs=1,
+                attrs=1,
+            ),
             defs.spam_returns_arg: new_var_counts(
                 posorkw=1,
             ),
@@ -902,6 +925,9 @@ def new_var_counts(*,
                 othercells=1,
                 purelocals=1,
             ),
+            defs.spam_annotated: new_var_counts(
+                posorkw=3,
+            ),
             defs.spam_full: new_var_counts(
                 posonly=2,
                 posorkw=2,

Lib/test/test_crossinterp.py

@@ -725,6 +725,39 @@ def test_user_exception(self):
         ])
 
 
+class CodeTests(_GetXIDataTests):
+
+    MODE = 'code'
+
+    def test_function_code(self):
+        self.assert_roundtrip_equal_not_identical([
+            *(f.__code__ for f in defs.FUNCTIONS),
+            *(f.__code__ for f in defs.FUNCTION_LIKE),
+        ])
+
+    def test_functions(self):
+        self.assert_not_shareable([
+            *defs.FUNCTIONS,
+            *defs.FUNCTION_LIKE,
+        ])
+
+    def test_other_objects(self):
+        self.assert_not_shareable([
+            None,
+            True,
+            False,
+            Ellipsis,
+            NotImplemented,
+            9999,
+            'spam',
+            b'spam',
+            (),
+            [],
+            {},
+            object(),
+        ])
+
+
 class ShareableTypeTests(_GetXIDataTests):
 
     MODE = 'xidata'
@@ -817,6 +850,13 @@ def test_object(self):
             object(),
         ])
 
+    def test_code(self):
+        # types.CodeType
+        self.assert_not_shareable([
+            *(f.__code__ for f in defs.FUNCTIONS),
+            *(f.__code__ for f in defs.FUNCTION_LIKE),
+        ])
+
     def test_function_object(self):
         for func in defs.FUNCTIONS:
             assert type(func) is types.FunctionType, func
@@ -935,12 +975,6 @@ def test_builtin_objects(self):
         self.assert_not_shareable([
             types.MappingProxyType({}),
             types.SimpleNamespace(),
-            # types.CodeType
-            defs.spam_minimal.__code__,
-            defs.spam_full.__code__,
-            defs.spam_CC.__code__,
-            defs.eggs_closure_C.__code__,
-            defs.ham_C_closure.__code__,
             # types.CellType
             types.CellType(),
             # types.FrameType

Modules/_testinternalcapi.c

@@ -1984,6 +1984,11 @@ get_crossinterp_data(PyObject *self, PyObject *args, PyObject *kwargs)
             goto error;
         }
     }
+    else if (strcmp(mode, "code") == 0) {
+        if (_PyCode_GetXIData(tstate, obj, xidata) != 0) {
+            goto error;
+        }
+    }
     else {
         PyErr_Format(PyExc_ValueError, "unsupported mode %R", modeobj);
         goto error;

Python/crossinterp_data_lookup.h

@@ -654,6 +654,30 @@ _tuple_shared(PyThreadState *tstate, PyObject *obj, _PyXIData_t *xidata)
     return -1;
 }
 
+// code
+
+PyObject *
+_PyCode_FromXIData(_PyXIData_t *xidata)
+{
+    return _PyMarshal_ReadObjectFromXIData(xidata);
+}
+
+int
+_PyCode_GetXIData(PyThreadState *tstate, PyObject *obj, _PyXIData_t *xidata)
+{
+    if (!PyCode_Check(obj)) {
+        _PyXIData_FormatNotShareableError(tstate, "expected code, got %R", obj);
+        return -1;
+    }
+    if (_PyMarshal_GetXIData(tstate, obj, xidata) < 0) {
+        return -1;
+    }
+    assert(_PyXIData_CHECK_NEW_OBJECT(xidata, _PyMarshal_ReadObjectFromXIData));
+    _PyXIData_SET_NEW_OBJECT(xidata, _PyCode_FromXIData);
+    return 0;
+}
+
+
 // registration
 
 static void