gh-132775: Add _PyBytes_GetXIData() (gh-133101)

This is the base for several other XIData wrappers, like pickle and marshal. It is essentially a refactor of the existing bytes XIData code.

Author: ericsnowcurrently

Include/internal/pycore_crossinterp.h

@@ -112,6 +112,8 @@ PyAPI_FUNC(void) _PyXIData_Clear(PyInterpreterState *, _PyXIData_t *);
     do { \
         (DATA)->free = (FUNC); \
     } while (0)
+#define _PyXIData_CHECK_FREE(DATA, FUNC) \
+    ((DATA)->free == (FUNC))
 // Additionally, some shareable types are essentially light wrappers
 // around other shareable types.  The xidatafunc of the wrapper
 // can often be implemented by calling the wrapped object's
@@ -123,6 +125,8 @@ PyAPI_FUNC(void) _PyXIData_Clear(PyInterpreterState *, _PyXIData_t *);
     do { \
         (DATA)->new_object = (FUNC); \
     } while (0)
+#define _PyXIData_CHECK_NEW_OBJECT(DATA, FUNC) \
+    ((DATA)->new_object == (FUNC))
 
 
 /* getting cross-interpreter data */
@@ -148,6 +152,25 @@ PyAPI_FUNC(int) _PyObject_GetXIData(
         PyObject *,
         _PyXIData_t *);
 
+// _PyObject_GetXIData() for bytes
+typedef struct {
+    const char *bytes;
+    Py_ssize_t len;
+} _PyBytes_data_t;
+PyAPI_FUNC(int) _PyBytes_GetData(PyObject *, _PyBytes_data_t *);
+PyAPI_FUNC(PyObject *) _PyBytes_FromData(_PyBytes_data_t *);
+PyAPI_FUNC(PyObject *) _PyBytes_FromXIData(_PyXIData_t *);
+PyAPI_FUNC(int) _PyBytes_GetXIData(
+        PyThreadState *,
+        PyObject *,
+        _PyXIData_t *);
+PyAPI_FUNC(_PyBytes_data_t *) _PyBytes_GetXIDataWrapped(
+        PyThreadState *,
+        PyObject *,
+        size_t,
+        xid_newobjfunc,
+        _PyXIData_t *);
+
 
 /* using cross-interpreter data */
 

Python/crossinterp.c

@@ -154,7 +154,7 @@ _PyXIData_InitWithSize(_PyXIData_t *xidata,
     // where it was allocated, so the interpreter is required.
     assert(interp != NULL);
     _PyXIData_Init(xidata, interp, NULL, obj, new_object);
-    xidata->data = PyMem_RawMalloc(size);
+    xidata->data = PyMem_RawCalloc(1, size);
     if (xidata->data == NULL) {
         return -1;
     }

Python/crossinterp_data_lookup.h

@@ -348,36 +348,98 @@ _PyXIData_UnregisterClass(PyThreadState *tstate, PyTypeObject *cls)
 
 // bytes
 
-struct _shared_bytes_data {
+int
+_PyBytes_GetData(PyObject *obj, _PyBytes_data_t *data)
+{
+    if (!PyBytes_Check(obj)) {
+        PyErr_Format(PyExc_TypeError, "expected bytes, got %R", obj);
+        return -1;
+    }
     char *bytes;
     Py_ssize_t len;
-};
+    if (PyBytes_AsStringAndSize(obj, &bytes, &len) < 0) {
+        return -1;
+    }
+    *data = (_PyBytes_data_t){
+        .bytes = bytes,
+        .len = len,
+    };
+    return 0;
+}
 
-static PyObject *
-_new_bytes_object(_PyXIData_t *xidata)
+PyObject *
+_PyBytes_FromData(_PyBytes_data_t *data)
 {
-    struct _shared_bytes_data *shared = (struct _shared_bytes_data *)(xidata->data);
-    return PyBytes_FromStringAndSize(shared->bytes, shared->len);
+    return PyBytes_FromStringAndSize(data->bytes, data->len);
+}
+
+PyObject *
+_PyBytes_FromXIData(_PyXIData_t *xidata)
+{
+    _PyBytes_data_t *data = (_PyBytes_data_t *)xidata->data;
+    assert(_PyXIData_OBJ(xidata) != NULL
+            && PyBytes_Check(_PyXIData_OBJ(xidata)));
+    return _PyBytes_FromData(data);
 }
 
 static int
-_bytes_shared(PyThreadState *tstate, PyObject *obj, _PyXIData_t *xidata)
+_bytes_shared(PyThreadState *tstate,
+              PyObject *obj, size_t size, xid_newobjfunc newfunc,
+              _PyXIData_t *xidata)
 {
+    assert(size >= sizeof(_PyBytes_data_t));
+    assert(newfunc != NULL);
     if (_PyXIData_InitWithSize(
-            xidata, tstate->interp, sizeof(struct _shared_bytes_data), obj,
-            _new_bytes_object
-            ) < 0)
+                        xidata, tstate->interp, size, obj, newfunc) < 0)
     {
         return -1;
     }
-    struct _shared_bytes_data *shared = (struct _shared_bytes_data *)xidata->data;
-    if (PyBytes_AsStringAndSize(obj, &shared->bytes, &shared->len) < 0) {
+    _PyBytes_data_t *data = (_PyBytes_data_t *)xidata->data;
+    if (_PyBytes_GetData(obj, data) < 0) {
         _PyXIData_Clear(tstate->interp, xidata);
         return -1;
     }
     return 0;
 }
 
+int
+_PyBytes_GetXIData(PyThreadState *tstate, PyObject *obj, _PyXIData_t *xidata)
+{
+    if (!PyBytes_Check(obj)) {
+        PyErr_Format(PyExc_TypeError, "expected bytes, got %R", obj);
+        return -1;
+    }
+    size_t size = sizeof(_PyBytes_data_t);
+    return _bytes_shared(tstate, obj, size, _PyBytes_FromXIData, xidata);
+}
+
+_PyBytes_data_t *
+_PyBytes_GetXIDataWrapped(PyThreadState *tstate,
+                          PyObject *obj, size_t size, xid_newobjfunc newfunc,
+                          _PyXIData_t *xidata)
+{
+    if (!PyBytes_Check(obj)) {
+        PyErr_Format(PyExc_TypeError, "expected bytes, got %R", obj);
+        return NULL;
+    }
+    if (size < sizeof(_PyBytes_data_t)) {
+        PyErr_Format(PyExc_ValueError, "expected size >= %d, got %d",
+                     sizeof(_PyBytes_data_t), size);
+        return NULL;
+    }
+    if (newfunc == NULL) {
+        if (size == sizeof(_PyBytes_data_t)) {
+            PyErr_SetString(PyExc_ValueError, "missing new_object func");
+            return NULL;
+        }
+        newfunc = _PyBytes_FromXIData;
+    }
+    if (_bytes_shared(tstate, obj, size, newfunc, xidata) < 0) {
+        return NULL;
+    }
+    return (_PyBytes_data_t *)xidata->data;
+}
+
 // str
 
 struct _shared_str_data {
@@ -608,7 +670,7 @@ _register_builtins_for_crossinterpreter_data(dlregistry_t *xidregistry)
     }
 
     // bytes
-    if (_xidregistry_add_type(xidregistry, &PyBytes_Type, _bytes_shared) != 0) {
+    if (_xidregistry_add_type(xidregistry, &PyBytes_Type, _PyBytes_GetXIData) != 0) {
         Py_FatalError("could not register bytes for cross-interpreter sharing");
     }