[pypy-commit] pypy py3k: Implement BytesIO.getbuffer()

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: py3k
Changeset: r61425:ac389af226dc
Date: 2013-02-18 19:30 +0100
http://bitbucket.org/pypy/pypy/changeset/ac389af226dc/

Log:	Implement BytesIO.getbuffer()

diff --git a/lib-python/3.2/test/test_memoryio.py b/lib-python/3.2/test/test_memoryio.py
--- a/lib-python/3.2/test/test_memoryio.py
+++ b/lib-python/3.2/test/test_memoryio.py
@@ -395,8 +395,11 @@
         self.assertEqual(bytes(buf), b"1234567890")
         # Trying to change the size of the BytesIO while a buffer is exported
         # raises a BufferError.
-        self.assertRaises(BufferError, memio.write, b'x' * 100)
-        self.assertRaises(BufferError, memio.truncate)
+        if support.check_impl_detail(pypy=False):
+            # PyPy export buffers differently, and allows reallocation
+            # of the underlying object.
+            self.assertRaises(BufferError, memio.write, b'x' * 100)
+            self.assertRaises(BufferError, memio.truncate)
         # Mutating the buffer updates the BytesIO
         buf[3:6] = b"abc"
         self.assertEqual(bytes(buf), b"123abc7890")
diff --git a/pypy/module/_io/interp_bytesio.py b/pypy/module/_io/interp_bytesio.py
--- a/pypy/module/_io/interp_bytesio.py
+++ b/pypy/module/_io/interp_bytesio.py
@@ -2,6 +2,7 @@
     TypeDef, generic_new_descr, GetSetProperty)
 from pypy.interpreter.gateway import interp2app, unwrap_spec
 from pypy.interpreter.error import OperationError, operationerrfmt
+from pypy.interpreter.buffer import RWBuffer
 from rpython.rlib.rarithmetic import r_longlong
 from pypy.module._io.interp_bufferedio import W_BufferedIOBase
 from pypy.module._io.interp_iobase import convert_size
@@ -14,6 +15,20 @@
         builder.append(buffer[i])
     return builder.build()
 
+class BytesIOBuffer(RWBuffer):
+    def __init__(self, w_bytesio):
+        self.w_bytesio = w_bytesio
+
+    def getlength(self):
+        return int(self.w_bytesio.string_size)
+
+    def getitem(self, index):
+        return self.w_bytesio.buf[index]
+
+    def setitem(self, index, char):
+        self.w_bytesio.buf[index] = char
+
+
 class W_BytesIO(W_BufferedIOBase):
     def __init__(self, space):
         W_BufferedIOBase.__init__(self, space)
@@ -123,6 +138,9 @@
 
         return space.wrap(size)
 
+    def getbuffer_w(self, space):
+        return space.wrap(BytesIOBuffer(self))
+
     def getvalue_w(self, space):
         self._check_closed(space)
         return space.wrapbytes(buffer2string(self.buf, 0, self.string_size))
@@ -213,6 +231,7 @@
     readinto = interp2app(W_BytesIO.readinto_w),
     write = interp2app(W_BytesIO.write_w),
     truncate = interp2app(W_BytesIO.truncate_w),
+    getbuffer = interp2app(W_BytesIO.getbuffer_w),
     getvalue = interp2app(W_BytesIO.getvalue_w),
     seek = interp2app(W_BytesIO.seek_w),
     tell = interp2app(W_BytesIO.tell_w),
diff --git a/pypy/module/_io/test/test_bytesio.py b/pypy/module/_io/test/test_bytesio.py
--- a/pypy/module/_io/test/test_bytesio.py
+++ b/pypy/module/_io/test/test_bytesio.py
@@ -75,3 +75,20 @@
         b = _io.BytesIO(b"hello")
         b.close()
         raises(ValueError, b.readinto, bytearray(b"hello"))
+
+    def test_getbuffer(self):
+        import _io
+        memio = _io.BytesIO(b"1234567890")
+        buf = memio.getbuffer()
+        assert bytes(buf) == b"1234567890"
+        memio.seek(5)
+        buf = memio.getbuffer()
+        assert bytes(buf) == b"1234567890"
+        # Mutating the buffer updates the BytesIO
+        buf[3:6] = b"abc"
+        assert bytes(buf) == b"123abc7890"
+        assert memio.getvalue() == b"123abc7890"
+        # After the buffer gets released, we can resize the BytesIO again
+        del buf
+        memio.truncate()
+