[pypy-commit] pypy default: initial support for numpy.frombuffer

[bdkearns]

Author: Brian Kearns <bdkearns at gmail.com>
Branch: 
Changeset: r71688:d7b779ef7782
Date: 2014-05-23 17:25 -0400
http://bitbucket.org/pypy/pypy/changeset/d7b779ef7782/

Log:	initial support for numpy.frombuffer

diff --git a/pypy/module/micronumpy/__init__.py b/pypy/module/micronumpy/__init__.py
--- a/pypy/module/micronumpy/__init__.py
+++ b/pypy/module/micronumpy/__init__.py
@@ -15,6 +15,7 @@
         'empty': 'ctors.zeros',
         'empty_like': 'ctors.empty_like',
         'fromstring': 'ctors.fromstring',
+        'frombuffer': 'ctors.frombuffer',
 
         'concatenate': 'arrayops.concatenate',
         'count_nonzero': 'arrayops.count_nonzero',
diff --git a/pypy/module/micronumpy/ctors.py b/pypy/module/micronumpy/ctors.py
--- a/pypy/module/micronumpy/ctors.py
+++ b/pypy/module/micronumpy/ctors.py
@@ -1,5 +1,6 @@
 from pypy.interpreter.error import OperationError, oefmt
 from pypy.interpreter.gateway import unwrap_spec, WrappedDefault
+from rpython.rlib.buffer import SubBuffer
 from rpython.rlib.rstring import strip_spaces
 from rpython.rtyper.lltypesystem import lltype, rffi
 from pypy.module.micronumpy import descriptor, loop
@@ -191,3 +192,54 @@
         return _fromstring_bin(space, s, count, length, dtype)
     else:
         return _fromstring_text(space, s, count, sep, length, dtype)
+
+
+def _getbuffer(space, w_buffer):
+    try:
+        return space.writebuf_w(w_buffer)
+    except OperationError as e:
+        if not e.match(space, space.w_TypeError):
+            raise
+        return space.readbuf_w(w_buffer)
+
+
+ at unwrap_spec(count=int, offset=int)
+def frombuffer(space, w_buffer, w_dtype=None, count=-1, offset=0):
+    dtype = space.interp_w(descriptor.W_Dtype,
+        space.call_function(space.gettypefor(descriptor.W_Dtype), w_dtype))
+    if dtype.elsize == 0:
+        raise oefmt(space.w_ValueError, "itemsize cannot be zero in type")
+
+    try:
+        buf = _getbuffer(space, w_buffer)
+    except OperationError as e:
+        if not e.match(space, space.w_TypeError):
+            raise
+        buf = _getbuffer(space, space.getattr(w_buffer, space.wrap('__buffer__')))
+
+    ts = buf.getlength()
+    if offset < 0 or offset > ts:
+        raise oefmt(space.w_ValueError,
+                    "offset must be non-negative and no greater than "
+                    "buffer length (%d)", ts)
+
+    s = ts - offset
+    if offset:
+        buf = SubBuffer(buf, offset, s)
+
+    n = count
+    itemsize = dtype.elsize
+    assert itemsize > 0
+    if n < 0:
+        if s % itemsize != 0:
+            raise oefmt(space.w_ValueError,
+                        "buffer size must be a multiple of element size")
+        n = s / itemsize
+    else:
+        if s < n * itemsize:
+            raise oefmt(space.w_ValueError,
+                        "buffer is smaller than requested size")
+
+    a = W_NDimArray.from_shape(space, [n], dtype=dtype)
+    loop.fromstring_loop(space, a, dtype, itemsize, buf.as_str())
+    return space.wrap(a)
diff --git a/pypy/module/micronumpy/test/test_ndarray.py b/pypy/module/micronumpy/test/test_ndarray.py
--- a/pypy/module/micronumpy/test/test_ndarray.py
+++ b/pypy/module/micronumpy/test/test_ndarray.py
@@ -3142,6 +3142,23 @@
         cls.w_float64val = cls.space.wrap(struct.pack('d', 300.4))
         cls.w_ulongval = cls.space.wrap(struct.pack('L', 12))
 
+    def test_frombuffer(self):
+        import numpy as np
+        exc = raises(AttributeError, np.frombuffer, None)
+        assert str(exc.value) == "'NoneType' object has no attribute '__buffer__'"
+        exc = raises(AttributeError, np.frombuffer, memoryview(self.data))
+        assert str(exc.value) == "'memoryview' object has no attribute '__buffer__'"
+        exc = raises(ValueError, np.frombuffer, self.data, 'S0')
+        assert str(exc.value) == "itemsize cannot be zero in type"
+        exc = raises(ValueError, np.frombuffer, self.data, offset=-1)
+        assert str(exc.value) == "offset must be non-negative and no greater than buffer length (32)"
+        exc = raises(ValueError, np.frombuffer, self.data, count=100)
+        assert str(exc.value) == "buffer is smaller than requested size"
+        for data in [self.data, buffer(self.data)]:
+            a = np.frombuffer(data)
+            for i in range(4):
+                assert a[i] == i + 1
+
     def test_fromstring(self):
         import sys
         from numpypy import fromstring, dtype