[pypy-commit] pypy PyBuffer: Add Buffer.w_getitem, fix memoryview.__getitem__ in the single-element case

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: PyBuffer
Changeset: r91029:278bd2e57510
Date: 2017-04-10 16:58 +0100
http://bitbucket.org/pypy/pypy/changeset/278bd2e57510/

Log:	Add Buffer.w_getitem, fix memoryview.__getitem__ in the single-
	element case

diff --git a/pypy/interpreter/buffer.py b/pypy/interpreter/buffer.py
--- a/pypy/interpreter/buffer.py
+++ b/pypy/interpreter/buffer.py
@@ -2,6 +2,7 @@
 from rpython.rlib.signature import signature
 from rpython.rlib import types
 
+
 class Buffer(object):
     """Abstract base class for buffers."""
     _attrs_ = ['readonly']
@@ -77,6 +78,20 @@
             if copiedbytes >= bytesize:
                 break
 
+    def w_getitem(self, space, idx):
+        from pypy.module.struct.formatiterator import UnpackFormatIterator
+        itemsize = self.getitemsize()
+        if itemsize == 1:
+            ch = self.as_binary()[idx]
+            return space.newint(ord(ch))
+        else:
+            # TODO: this probably isn't very fast
+            buf = SubBuffer(self.as_binary(), idx, itemsize)
+            fmtiter = UnpackFormatIterator(space, buf)
+            fmtiter.length = buf.getlength()
+            fmtiter.interpret(self.getformat())
+            return fmtiter.result_w[0]
+
 
 class SimpleBuffer(Buffer):
     _attrs_ = ['readonly', 'data']
@@ -113,6 +128,10 @@
     def getstrides(self):
         return [1]
 
+    def w_getitem(self, space, idx):
+        ch = self.data[idx]
+        return space.newint(ord(ch))
+
 
 class BinaryBuffer(object):
     """Base class for buffers of bytes"""
diff --git a/pypy/objspace/std/memoryobject.py b/pypy/objspace/std/memoryobject.py
--- a/pypy/objspace/std/memoryobject.py
+++ b/pypy/objspace/std/memoryobject.py
@@ -223,16 +223,7 @@
                 shape = self.getshape()
                 strides = self.getstrides()
                 idx = self.lookup_dimension(space, shape, strides, 0, 0, start)
-                if itemsize == 1:
-                    ch = self.buf.getitem(idx)
-                    return space.newint(ord(ch))
-                else:
-                    # TODO: this probably isn't very fast
-                    buf = SubBuffer(self.buf.as_binary(), idx, itemsize)
-                    fmtiter = UnpackFormatIterator(space, buf)
-                    fmtiter.length = buf.getlength()
-                    fmtiter.interpret(self.buf.getformat())
-                    return fmtiter.result_w[0]
+                return self.buf.w_getitem(space, idx)
             else:
                 raise oefmt(space.w_NotImplementedError, "multi-dimensional sub-views are not implemented")
         elif is_slice: