Merge pull request #17 from jakirkham/add_getbuffer

Add getbuffer

src/cybuffer.pxd

@@ -1,3 +1,10 @@
+include "config.pxi"
+
+
+IF PY2K:
+    cpdef getbuffer(obj, Py_ssize_t offset=*, Py_ssize_t size=*)
+
+
 cdef class cybuffer(object):
     cdef readonly object obj
 

src/cybuffer.pyx

@@ -11,7 +11,6 @@ cimport cpython.buffer
 cimport cpython.bytes
 cimport cpython.list
 cimport cpython.mem
-cimport cpython.oldbuffer
 cimport cpython.tuple
 
 from cpython.array cimport array
@@ -55,6 +54,31 @@ cdef extern from *:
     void PyTuple_SET_ITEM_INC(object, Py_ssize_t, object)
 
 
+IF PY2K:
+    cimport cpython.oldbuffer
+    from cpython.oldbuffer cimport (
+        PyBuffer_FromReadWriteObject, PyBuffer_FromObject
+    )
+
+
+    cdef extern from "Python.h":
+        int PyObject_AsWriteBuffer(object obj,
+                                   void **buffer,
+                                   Py_ssize_t *buffer_len) except -1
+
+
+    cpdef getbuffer(obj, Py_ssize_t offset=0, Py_ssize_t size=-1):
+        cdef void* buf_ptr
+        cdef Py_ssize_t buf_len
+
+        try:
+            PyObject_AsWriteBuffer(obj, &buf_ptr, &buf_len)
+        except TypeError:
+            return PyBuffer_FromObject(obj, offset, size)
+        else:
+            return PyBuffer_FromReadWriteObject(obj, offset, size)
+
+
 cdef tuple pointer_to_tuple(int n, Py_ssize_t* p):
     cdef int i
     cdef object p_i
@@ -130,13 +154,8 @@ cdef class cybuffer(object):
         self.obj = data
 
         # Fallback to old buffer protocol on Python 2 if necessary
-        if PY2K and not cpython.buffer.PyObject_CheckBuffer(self.obj):
-            try:
-                data = cpython.oldbuffer.PyBuffer_FromReadWriteObject(
-                    self.obj, 0, -1
-                )
-            except TypeError:
-                data = cpython.oldbuffer.PyBuffer_FromObject(self.obj, 0, -1)
+        if PY2K and not cpython.buffer.PyObject_CheckBuffer(data):
+            data = getbuffer(data)
 
         # Fill out our buffer based on the data
         cpython.buffer.PyObject_GetBuffer(data, &self._buf, PyBUF_FULL_RO)

tests/test_cybuffer.py

@@ -16,14 +16,37 @@
 
 
 try:
-    buffer
-except NameError:
-    buffer = memoryview
+    from cybuffer import getbuffer
+except ImportError:
+    getbuffer = memoryview
 
 
 Py_UNICODE_SIZE = array.array('u').itemsize
 
 
+@pytest.mark.skipif(
+    sys.version_info[0] != 2, reason="getbuffer is Python 2 only"
+)
+@pytest.mark.parametrize("v, to_char", [
+    (b"abcdefghi", lambda c: c),
+    (bytearray(b"abcdefghi"), chr),
+])
+def test_getbuffer(v, to_char):
+    # Initialize buffers
+    b = getbuffer(v)
+    m = memoryview(v)
+    mb = memoryview(b)
+
+    # Validate type
+    assert isinstance(b, buffer)
+
+    # Validate content
+    assert list(b) == list(map(to_char, v))
+
+    # Validate permissions
+    assert mb.readonly == m.readonly
+
+
 def test_empty_constructor():
     with pytest.raises(TypeError):
         b = cybuffer()
@@ -93,7 +116,7 @@ def test_1d_arrays(f):
     # Initialize buffers
     v = array.array(f, [0, 1, 2, 3, 4])
     b = cybuffer(v)
-    m = memoryview(buffer(v))
+    m = memoryview(getbuffer(v))
 
     # Validate format
     assert b.format == v.typecode
@@ -105,10 +128,7 @@ def test_1d_arrays(f):
     assert b.contiguous
 
     # Validate permissions
-    if isinstance(b, memoryview):
-        assert b.readonly
-    else:
-        assert not b.readonly
+    assert not b.readonly
 
     # Test methods
     assert b.tolist() == v.tolist()
@@ -128,7 +148,7 @@ def test_1d_text_arrays(f, s):
     # Initialize buffers
     v = array.array(f, s)
     b = cybuffer(v)
-    m = memoryview(buffer(v))
+    m = memoryview(getbuffer(v))
 
     # Validate format
     assert b.itemsize == v.itemsize
@@ -145,10 +165,7 @@ def test_1d_text_arrays(f, s):
     assert b.contiguous
 
     # Validate permissions
-    if isinstance(b, memoryview):
-        assert b.readonly
-    else:
-        assert not b.readonly
+    assert not b.readonly
 
     # Test methods
     assert b.tolist() == list(map(ord, v))
@@ -160,7 +177,7 @@ def test_mmap():
     with contextlib.closing(mmap.mmap(-1, 10, access=mmap.ACCESS_WRITE)) as v:
         # Initialize buffers
         b = cybuffer(v)
-        m = memoryview(buffer(v))
+        m = memoryview(getbuffer(v))
 
         # Validate format
         assert b.format == m.format