Add support for native PickleBuffer as_buffer

Author: qunaibit

graalpython/com.oracle.graal.python.cext/src/capi.c

@@ -96,6 +96,14 @@ typedef struct {
     PyObject *mapping;
 } mappingproxyobject;
 
+typedef struct {
+    PyObject_HEAD
+    /* The view exported by the original object */
+    Py_buffer view;
+    PyObject *weakreflist;
+} PyPickleBufferObject;
+
+
 static void object_dealloc(PyObject *self) {
     Py_TYPE(self)->tp_free(self);
 }
@@ -278,6 +286,16 @@ void bytearray_releasebuffer(PyByteArrayObject *obj, Py_buffer *view);
 int memory_getbuf(PyMemoryViewObject *self, Py_buffer *view, int flags);
 void memory_releasebuf(PyMemoryViewObject *self, Py_buffer *view);
 
+/* PICKLEBUFFER */
+static int
+picklebuf_getbuf(PyPickleBufferObject *self, Py_buffer *view, int flags)
+{
+    PyObject *self_view_obj = GraalPyTruffle_PickleBuffer_viewobj(self);
+    return PyObject_GetBuffer(self_view_obj, view, flags);
+}
+
+static void empty_releasebuf(PyObject *self, Py_buffer *view) {}
+
 static void initialize_bufferprocs() {
     static PyBufferProcs bytes_as_buffer = {
         (getbufferproc)bytes_buffer_getbuffer,       /* bf_getbuffer */
@@ -305,6 +323,12 @@ static void initialize_bufferprocs() {
     array_as_buffer.bf_releasebuffer = GraalPyTruffle_Array_releasebuffer,
     Arraytype.tp_as_buffer = &array_as_buffer;
     GraalPy_set_PyTypeObject_tp_as_buffer(&Arraytype, &array_as_buffer);
+
+    static PyBufferProcs picklebuf_as_buffer;
+    picklebuf_as_buffer.bf_getbuffer = (getbufferproc)picklebuf_getbuf,
+    picklebuf_as_buffer.bf_releasebuffer = empty_releasebuf,
+    PyPickleBuffer_Type.tp_as_buffer = &picklebuf_as_buffer;
+    GraalPy_set_PyTypeObject_tp_as_buffer(&PyPickleBuffer_Type, &picklebuf_as_buffer);
 }
 
 int is_builtin_type(PyTypeObject *tp) {

graalpython/com.oracle.graal.python.cext/src/capi.h

@@ -342,6 +342,8 @@ PY_TRUFFLE_TYPE(UnionType_Type, 		"_ctypes.UnionType",		&PyType_Type, sizeof(PyO
 PY_TRUFFLE_TYPE(PyCPointerType_Type,	"PyCPointerType", 		    &PyType_Type, sizeof(PyObject)) \
 PY_TRUFFLE_TYPE(PyCArrayType_Type,		"PyCArrayType", 			&PyType_Type, sizeof(PyObject)) \
 PY_TRUFFLE_TYPE(PyCoro_Type, 			"coroutine", 				&PyType_Type, sizeof(PyCoroObject)) \
+/* PyPickleBufferObject (PyObject_HEAD + Py_buffer + PyObject*) is defined within Objects/picklebufobject.c, so its not exposed. */ \
+PY_TRUFFLE_TYPE(PyPickleBuffer_Type, 	"_pickle.PickleBuffer",     &PyType_Type, sizeof(PyPickleBufferObject)) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(_PyAIterWrapper_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(_PyAsyncGenASend_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(_PyAsyncGenAThrow_Type) \
@@ -385,7 +387,6 @@ PY_TRUFFLE_TYPE_UNIMPLEMENTED(PyODictItems_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(PyODictIter_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(PyODictKeys_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(PyODictValues_Type) \
-PY_TRUFFLE_TYPE_UNIMPLEMENTED(PyPickleBuffer_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(PyRangeIter_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(PySetIter_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(PySortWrapper_Type) \

graalpython/com.oracle.graal.python.cext/src/obmalloc.c

@@ -110,7 +110,7 @@ int PyTruffle_AllocMemory(size_t size) {
 
 void PyTruffle_FreeMemory(size_t size) {
     if (PyTruffle_AllocatedMemory < size) {
-        PyTruffle_Log(PY_TRUFFLE_LOG_INFO, "PyTruffle_FreeMemory: freed memory size (%lu) is larger than allocated memory size (%lu)\n", size, PyTruffle_AllocMemory);
+        PyTruffle_Log(PY_TRUFFLE_LOG_INFO, "PyTruffle_FreeMemory: freed memory size (%lu) is larger than allocated memory size (%lu)\n", size, PyTruffle_AllocatedMemory);
         PyTruffle_AllocatedMemory = size;
     }
     PyTruffle_AllocatedMemory -= size;

graalpython/com.oracle.graal.python.test/src/tests/cpyext/test_bytes.py

@@ -582,6 +582,45 @@ def test_create_from_buffer_exception(self):
         self.assertRaises(ValueError, bytes, TestType())
         self.assertRaises(ValueError, bytearray, TestType())
 
+    def test_tp_as_buffer_pickle(self):
+        TestAsBufferPickle = CPyExtType(
+            "TestAsBufferPickle",
+            """
+            static PyObject* get_bytes(PyObject* self, PyObject *obj) {
+                Py_buffer view;
+                if (PyObject_GetBuffer(obj, &view, PyBUF_FULL_RO) != 0)
+                    return NULL;
+                PyObject* bytes = PyBytes_FromStringAndSize(view.buf, view.len);
+                PyBuffer_Release(&view);
+                return bytes;
+            }
+            """,
+            tp_methods='{"get_bytes", (PyCFunction)get_bytes, METH_O | METH_CLASS, ""}',
+        )
+        TestType = CPyExtType(
+            "TestMemoryViewBufferPickle",
+            """
+            char buf[] = {1,2,3,4};
+            int getbuffer(TestMemoryViewBufferPickleObject *self, Py_buffer *view, int flags) {
+                return PyBuffer_FillInfo(view, (PyObject*)self, buf, 4, 1, flags);
+            }
+            void releasebuffer(TestMemoryViewBufferPickleObject *self, Py_buffer *view) {}
+            
+            static PyBufferProcs as_buffer = {
+                (getbufferproc)getbuffer,
+                (releasebufferproc)releasebuffer,
+            };
+            """,
+            tp_as_buffer='&as_buffer',
+        )
+        obj = TestType()
+        b = bytes([1,2,3,4]) # same as `buf[] = {1,2,3,4};`
+        import pickle
+        b2 = TestAsBufferPickle.get_bytes(pickle.PickleBuffer(memoryview(obj)))
+        b1 = TestAsBufferPickle.get_bytes(pickle.PickleBuffer(bytearray([1,2,3,4])))
+        assert b == b1
+        assert b == b2
+
 
 class TestNativeSubclass(unittest.TestCase):
     def test_builtins(self):

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/cext/PythonCextPickleBufferBuiltins.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) 2023, 2024, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * The Universal Permissive License (UPL), Version 1.0
+ *
+ * Subject to the condition set forth below, permission is hereby granted to any
+ * person obtaining a copy of this software, associated documentation and/or
+ * data (collectively the "Software"), free of charge and under any and all
+ * copyright rights in the Software, and any and all patent rights owned or
+ * freely licensable by each licensor hereunder covering either (i) the
+ * unmodified Software as contributed to or provided by such licensor, or (ii)
+ * the Larger Works (as defined below), to deal in both
+ *
+ * (a) the Software, and
+ *
+ * (b) any piece of software and/or hardware listed in the lrgrwrks.txt file if
+ * one is included with the Software each a "Larger Work" to which the Software
+ * is contributed by such licensors),
+ *
+ * without restriction, including without limitation the rights to copy, create
+ * derivative works of, display, perform, and distribute the Software and make,
+ * use, sell, offer for sale, import, export, have made, and have sold the
+ * Software and the Larger Work(s), and to sublicense the foregoing rights on
+ * either these or other terms.
+ *
+ * This license is subject to the following condition:
+ *
+ * The above copyright notice and either this complete permission notice or at a
+ * minimum a reference to the UPL must be included in all copies or substantial
+ * portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+package com.oracle.graal.python.builtins.modules.cext;
+
+import static com.oracle.graal.python.builtins.modules.cext.PythonCextBuiltins.CApiCallPath.Ignored;
+import static com.oracle.graal.python.builtins.objects.cext.capi.transitions.ArgDescriptor.PyObject;
+import static com.oracle.graal.python.builtins.objects.cext.capi.transitions.ArgDescriptor.PyObjectBorrowed;
+import static com.oracle.graal.python.runtime.exception.PythonErrorType.ValueError;
+
+import com.oracle.graal.python.builtins.modules.cext.PythonCextBuiltins.CApiBuiltin;
+import com.oracle.graal.python.builtins.modules.cext.PythonCextBuiltins.CApiUnaryBuiltinNode;
+import com.oracle.graal.python.builtins.modules.pickle.PPickleBuffer;
+import com.oracle.graal.python.builtins.objects.buffer.PythonBufferAccessLibrary;
+import com.oracle.graal.python.nodes.ErrorMessages;
+import com.oracle.graal.python.nodes.PRaiseNode;
+import com.oracle.truffle.api.dsl.Bind;
+import com.oracle.truffle.api.dsl.Cached;
+import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.library.CachedLibrary;
+import com.oracle.truffle.api.nodes.Node;
+
+public final class PythonCextPickleBufferBuiltins {
+
+    @CApiBuiltin(ret = PyObjectBorrowed, args = {PyObject}, call = Ignored)
+    abstract static class PyTruffle_PickleBuffer_viewobj extends CApiUnaryBuiltinNode {
+        @Specialization
+        static Object getviewobj(PPickleBuffer object,
+                        @Bind("this") Node inliningTarget,
+                        @CachedLibrary(limit = "3") PythonBufferAccessLibrary bufferLib,
+                        @Cached PRaiseNode.Lazy raiseNode) {
+            Object owner = null;
+            if (object.getView() != null) {
+                owner = bufferLib.getOwner(object.getView());
+            }
+            if (owner == null) {
+                throw raiseNode.get(inliningTarget).raise(ValueError, ErrorMessages.OP_FORBIDDEN_ON_OBJECT, "PickleBuffer");
+            }
+            return owner;
+        }
+
+    }
+}