gh-102500: Implement PEP 688 (#102521)

Co-authored-by: Kumar Aditya <59607654+kumaraditya303@users.noreply.github.com>

Author: JelleZijlstra

Include/internal/pycore_global_objects_fini_generated.h

@@ -81,6 +81,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(__await__)
         STRUCT_FOR_ID(__bases__)
         STRUCT_FOR_ID(__bool__)
+        STRUCT_FOR_ID(__buffer__)
         STRUCT_FOR_ID(__build_class__)
         STRUCT_FOR_ID(__builtins__)
         STRUCT_FOR_ID(__bytes__)
@@ -180,6 +181,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(__rdivmod__)
         STRUCT_FOR_ID(__reduce__)
         STRUCT_FOR_ID(__reduce_ex__)
+        STRUCT_FOR_ID(__release_buffer__)
         STRUCT_FOR_ID(__repr__)
         STRUCT_FOR_ID(__reversed__)
         STRUCT_FOR_ID(__rfloordiv__)
@@ -610,6 +612,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(reducer_override)
         STRUCT_FOR_ID(registry)
         STRUCT_FOR_ID(rel_tol)
+        STRUCT_FOR_ID(release)
         STRUCT_FOR_ID(reload)
         STRUCT_FOR_ID(repl)
         STRUCT_FOR_ID(replace)

Include/internal/pycore_global_strings.h

@@ -0,0 +1,17 @@
+#ifndef Py_INTERNAL_MEMORYOBJECT_H
+#define Py_INTERNAL_MEMORYOBJECT_H
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+PyObject *
+PyMemoryView_FromObjectAndFlags(PyObject *v, int flags);
+
+#ifdef __cplusplus
+}
+#endif
+#endif /* !Py_INTERNAL_MEMORYOBJECT_H */

Include/internal/pycore_memoryobject.h

@@ -138,6 +138,8 @@ _Py_type_getattro(PyTypeObject *type, PyObject *name);
 PyObject *_Py_slot_tp_getattro(PyObject *self, PyObject *name);
 PyObject *_Py_slot_tp_getattr_hook(PyObject *self, PyObject *name);
 
+PyAPI_DATA(PyTypeObject) _PyBufferWrapper_Type;
+
 PyObject *
 _PySuper_Lookup(PyTypeObject *su_type, PyObject *su_obj, PyObject *name, int *meth_found);
 PyObject *

Include/internal/pycore_runtime_init_generated.h

@@ -104,7 +104,7 @@ PyAPI_FUNC(void) PyBuffer_Release(Py_buffer *view);
 /* Maximum number of dimensions */
 #define PyBUF_MAX_NDIM 64
 
-/* Flags for getting buffers */
+/* Flags for getting buffers. Keep these in sync with inspect.BufferFlags. */
 #define PyBUF_SIMPLE 0
 #define PyBUF_WRITABLE 0x0001
 

Include/internal/pycore_typeobject.h

@@ -49,7 +49,7 @@ def _f(): pass
            "Mapping", "MutableMapping",
            "MappingView", "KeysView", "ItemsView", "ValuesView",
            "Sequence", "MutableSequence",
-           "ByteString",
+           "ByteString", "Buffer",
            ]
 
 # This module has been renamed from collections.abc to _collections_abc to
@@ -439,6 +439,21 @@ def __subclasshook__(cls, C):
         return NotImplemented
 
 
+class Buffer(metaclass=ABCMeta):
+
+    __slots__ = ()
+
+    @abstractmethod
+    def __buffer__(self, flags: int, /) -> memoryview:
+        raise NotImplementedError
+
+    @classmethod
+    def __subclasshook__(cls, C):
+        if cls is Buffer:
+            return _check_methods(C, "__buffer__")
+        return NotImplemented
+
+
 class _CallableGenericAlias(GenericAlias):
     """ Represent `Callable[argtypes, resulttype]`.
 

Include/internal/pycore_unicodeobject_generated.h

@@ -43,6 +43,7 @@
     "Attribute",
     "BlockFinder",
     "BoundArguments",
+    "BufferFlags",
     "CORO_CLOSED",
     "CORO_CREATED",
     "CORO_RUNNING",
@@ -3312,6 +3313,28 @@ def signature(obj, *, follow_wrapped=True, globals=None, locals=None, eval_str=F
                                    globals=globals, locals=locals, eval_str=eval_str)
 
 
+class BufferFlags(enum.IntFlag):
+    SIMPLE = 0x0
+    WRITABLE = 0x1
+    FORMAT = 0x4
+    ND = 0x8
+    STRIDES = 0x10 | ND
+    C_CONTIGUOUS = 0x20 | STRIDES
+    F_CONTIGUOUS = 0x40 | STRIDES
+    ANY_CONTIGUOUS = 0x80 | STRIDES
+    INDIRECT = 0x100 | STRIDES
+    CONTIG = ND | WRITABLE
+    CONTIG_RO = ND
+    STRIDED = STRIDES | WRITABLE
+    STRIDED_RO = STRIDES
+    RECORDS = STRIDES | WRITABLE | FORMAT
+    RECORDS_RO = STRIDES | FORMAT
+    FULL = INDIRECT | WRITABLE | FORMAT
+    FULL_RO = INDIRECT | FORMAT
+    READ = 0x100
+    WRITE = 0x200
+
+
 def _main():
     """ Logic for inspecting an object given at command line """
     import argparse

Include/pybuffer.h

@@ -17,6 +17,7 @@
 import unittest
 from test import support
 from test.support import os_helper
+import inspect
 from itertools import permutations, product
 from random import randrange, sample, choice
 import warnings
@@ -4438,5 +4439,146 @@ def test_pybuffer_size_from_format(self):
                              struct.calcsize(format))
 
 
+class TestPythonBufferProtocol(unittest.TestCase):
+    def test_basic(self):
+        class MyBuffer:
+            def __buffer__(self, flags):
+                return memoryview(b"hello")
+
+        mv = memoryview(MyBuffer())
+        self.assertEqual(mv.tobytes(), b"hello")
+        self.assertEqual(bytes(MyBuffer()), b"hello")
+
+    def test_bad_buffer_method(self):
+        class MustReturnMV:
+            def __buffer__(self, flags):
+                return 42
+
+        self.assertRaises(TypeError, memoryview, MustReturnMV())
+
+        class NoBytesEither:
+            def __buffer__(self, flags):
+                return b"hello"
+
+        self.assertRaises(TypeError, memoryview, NoBytesEither())
+
+        class WrongArity:
+            def __buffer__(self):
+                return memoryview(b"hello")
+
+        self.assertRaises(TypeError, memoryview, WrongArity())
+
+    def test_release_buffer(self):
+        class WhatToRelease:
+            def __init__(self):
+                self.held = False
+                self.ba = bytearray(b"hello")
+
+            def __buffer__(self, flags):
+                if self.held:
+                    raise TypeError("already held")
+                self.held = True
+                return memoryview(self.ba)
+
+            def __release_buffer__(self, buffer):
+                self.held = False
+
+        wr = WhatToRelease()
+        self.assertFalse(wr.held)
+        with memoryview(wr) as mv:
+            self.assertTrue(wr.held)
+            self.assertEqual(mv.tobytes(), b"hello")
+        self.assertFalse(wr.held)
+
+    def test_same_buffer_returned(self):
+        class WhatToRelease:
+            def __init__(self):
+                self.held = False
+                self.ba = bytearray(b"hello")
+                self.created_mv = None
+
+            def __buffer__(self, flags):
+                if self.held:
+                    raise TypeError("already held")
+                self.held = True
+                self.created_mv = memoryview(self.ba)
+                return self.created_mv
+
+            def __release_buffer__(self, buffer):
+                assert buffer is self.created_mv
+                self.held = False
+
+        wr = WhatToRelease()
+        self.assertFalse(wr.held)
+        with memoryview(wr) as mv:
+            self.assertTrue(wr.held)
+            self.assertEqual(mv.tobytes(), b"hello")
+        self.assertFalse(wr.held)
+
+    def test_buffer_flags(self):
+        class PossiblyMutable:
+            def __init__(self, data, mutable) -> None:
+                self._data = bytearray(data)
+                self._mutable = mutable
+
+            def __buffer__(self, flags):
+                if flags & inspect.BufferFlags.WRITABLE:
+                    if not self._mutable:
+                        raise RuntimeError("not mutable")
+                    return memoryview(self._data)
+                else:
+                    return memoryview(bytes(self._data))
+
+        mutable = PossiblyMutable(b"hello", True)
+        immutable = PossiblyMutable(b"hello", False)
+        with memoryview._from_flags(mutable, inspect.BufferFlags.WRITABLE) as mv:
+            self.assertEqual(mv.tobytes(), b"hello")
+            mv[0] = ord(b'x')
+            self.assertEqual(mv.tobytes(), b"xello")
+        with memoryview._from_flags(mutable, inspect.BufferFlags.SIMPLE) as mv:
+            self.assertEqual(mv.tobytes(), b"xello")
+            with self.assertRaises(TypeError):
+                mv[0] = ord(b'h')
+            self.assertEqual(mv.tobytes(), b"xello")
+        with memoryview._from_flags(immutable, inspect.BufferFlags.SIMPLE) as mv:
+            self.assertEqual(mv.tobytes(), b"hello")
+            with self.assertRaises(TypeError):
+                mv[0] = ord(b'x')
+            self.assertEqual(mv.tobytes(), b"hello")
+
+        with self.assertRaises(RuntimeError):
+            memoryview._from_flags(immutable, inspect.BufferFlags.WRITABLE)
+        with memoryview(immutable) as mv:
+            self.assertEqual(mv.tobytes(), b"hello")
+            with self.assertRaises(TypeError):
+                mv[0] = ord(b'x')
+            self.assertEqual(mv.tobytes(), b"hello")
+
+    def test_call_builtins(self):
+        ba = bytearray(b"hello")
+        mv = ba.__buffer__(0)
+        self.assertEqual(mv.tobytes(), b"hello")
+        ba.__release_buffer__(mv)
+        with self.assertRaises(OverflowError):
+            ba.__buffer__(sys.maxsize + 1)
+
+    @unittest.skipIf(_testcapi is None, "requires _testcapi")
+    def test_c_buffer(self):
+        buf = _testcapi.testBuf()
+        self.assertEqual(buf.references, 0)
+        mv = buf.__buffer__(0)
+        self.assertIsInstance(mv, memoryview)
+        self.assertEqual(mv.tobytes(), b"test")
+        self.assertEqual(buf.references, 1)
+        buf.__release_buffer__(mv)
+        self.assertEqual(buf.references, 0)
+        with self.assertRaises(ValueError):
+            mv.tobytes()
+        # Calling it again doesn't cause issues
+        with self.assertRaises(ValueError):
+            buf.__release_buffer__(mv)
+        self.assertEqual(buf.references, 0)
+
+
 if __name__ == "__main__":
     unittest.main()