[pypy-commit] pypy numpy-data-buffer: implement numpy.frombuffer

[timo_jbo]

Author: Timo Paulssen <timonator at perpetuum-immobile.de>
Branch: numpy-data-buffer
Changeset: r47793:f30d7a06caf9
Date: 2011-10-03 16:23 +0200
http://bitbucket.org/pypy/pypy/changeset/f30d7a06caf9/

Log:	implement numpy.frombuffer

diff --git a/lib_pypy/numpy/__init__.py b/lib_pypy/numpy/__init__.py
--- a/lib_pypy/numpy/__init__.py
+++ b/lib_pypy/numpy/__init__.py
@@ -86,24 +86,34 @@
 
     return array(result)
 
+def __from_buffer_or_datastring(buf_or_str, dt, count, offset=0):
+    _dtype = dtype(dt)
+
+    if count > 0:
+        length = count * _dtype.itemsize
+        if length + offset > len(buf_or_str):
+            raise ValueError("length of string (%d) not enough for %d %s" %
+                             (len(buf_or_str), count, _dtype))
+
+        buf_or_str = buf_or_str[offset:length+offset]
+    else:
+        length = len(buf_or_str) - offset
+        buf_or_str = buf_or_str[offset:]
+        if len(buf_or_str) % _dtype.itemsize != 0:
+            raise ValueError("length of string (%d) not evenly dividable by size of dtype (%d)" %
+                             (len(buf_or_str), _dtype.itemsize))
+
+    arr = empty(length / _dtype.itemsize, dtype=_dtype)
+    arr.data[:length] = buf_or_str
+
+    return arr
+
+def frombuffer(buf, dtype=float, count=-1, offset=0):
+    return __from_buffer_or_datastring(buf, dtype, count, offset)
+
 def fromstring(s, dtype=float, count=-1, sep=''):
-    from _numpy import dtype as dt
     if sep:
         raise NotImplementedError("Cannot use fromstring with a separator yet")
 
-    _dtype = dt(dtype)
-    if count > 0:
-        length = count * _dtype.itemsize
-        if length > len(s):
-            raise ValueError("length of string (%d) not enough for %d %s" % (len(s), count, _dtype))
-        s = s[:length]
-    else:
-        length = len(s)
-        if len(s) % _dtype.itemsize != 0:
-            raise ValueError("length of string (%d) not evenly dividable by size of dtype (%d)" % (len(s), _dtype.itemsize))
+    return __from_buffer_or_datastring(s, dtype, count)
 
-    arr = empty(length / _dtype.itemsize, dtype=_dtype)
-    print len(arr.data), len(s), length, _dtype.itemsize
-    arr.data[:length] = s
-
-    return arr
diff --git a/lib_pypy/pypy_test/test_numpy.py b/lib_pypy/pypy_test/test_numpy.py
--- a/lib_pypy/pypy_test/test_numpy.py
+++ b/lib_pypy/pypy_test/test_numpy.py
@@ -96,3 +96,24 @@
         assert len(a) == 4
         for i in range(4):
             assert a[i] == i
+
+    def test_frombuffer(self):
+        from numpy import frombuffer
+        import struct
+
+        data = "\0\1\2\3\4"
+        a = frombuffer(data, dtype="int8", count=3)
+        b = frombuffer(data, dtype="int8", count=3, offset=1)
+        # does this depend on the machine architecture? byte-order?
+        assert a[0] == 0
+        assert b[0] == a[1] == 1
+
+        data = struct.pack("iiii", 0, 1, 0, 0)
+        a = frombuffer(data, dtype="i", count=1)[0]
+        b = frombuffer(data, dtype="i", count=1, offset=1)[0]
+        c = frombuffer(data, dtype="i", count=1, offset=2)[0]
+        d = frombuffer(data, dtype="i", count=1, offset=3)[0]
+        assert a == 0
+        assert b == 1 << 24
+        assert c == 1 << 16
+        assert d == 1 << 8