Repository: numba/numba

Issue:
Support for numpy.frombuffer() in nopython mode
It would be very handy to support `numpy.frombuffer()` in nopython mode.  Then buffers could be easily recast to take advantage of the rest of the Numpy API.



Discussion:

Resolving patch:
diff --git a/numba/pythonapi.py b/numba/pythonapi.py
--- a/numba/pythonapi.py
+++ b/numba/pythonapi.py
@@ -834,17 +834,18 @@ def nrt_adapt_ndarray_to_python(self, aryty, ary, dtypeptr):
         assert self.context.enable_nrt, "NRT required"
 
         intty = ir.IntType(32)
-        fnty = Type.function(self.pyobj, [self.voidptr, intty, self.pyobj])
+        fnty = Type.function(self.pyobj,
+                             [self.voidptr, intty, intty, self.pyobj])
         fn = self._get_function(fnty, name="NRT_adapt_ndarray_to_python")
         fn.args[0].add_attribute(lc.ATTR_NO_CAPTURE)
-        dtype = numpy_support.as_dtype(aryty.dtype)
 
         ndim = self.context.get_constant(types.int32, aryty.ndim)
+        writable = self.context.get_constant(types.int32, int(aryty.mutable))
 
         aryptr = cgutils.alloca_once_value(self.builder, ary)
         return self.builder.call(fn, [self.builder.bitcast(aryptr,
                                                            self.voidptr),
-                                      ndim, dtypeptr])
+                                      ndim, writable, dtypeptr])
 
     def nrt_adapt_ndarray_from_python(self, ary, ptr):
         assert self.context.enable_nrt
diff --git a/numba/targets/arrayobj.py b/numba/targets/arrayobj.py
--- a/numba/targets/arrayobj.py
+++ b/numba/targets/arrayobj.py
@@ -90,7 +90,7 @@ def shape(self):
 
 def get_itemsize(context, array_type):
     """
-    Return the item size for the given array type.
+    Return the item size for the given array or buffer type.
     """
     llty = context.get_data_type(array_type.dtype)
     return context.get_abi_sizeof(llty)
@@ -2058,3 +2058,43 @@ def array_copy(context, builder, sig, args):
             builder.store(builder.load(src_ptr), dest_ptr)
 
     return ret._getvalue()
+
+
+@builtin
+@implement(numpy.frombuffer, types.Kind(types.Buffer))
+@implement(numpy.frombuffer, types.Kind(types.Buffer), types.Kind(types.DTypeSpec))
+def np_frombuffer(context, builder, sig, args):
+    bufty = sig.args[0]
+    aryty = sig.return_type
+
+    buf = make_array(bufty)(context, builder, value=args[0])
+    out_ary_ty = make_array(aryty)
+    out_ary = out_ary_ty(context, builder)
+    out_datamodel = out_ary._datamodel
+
+    itemsize = get_itemsize(context, aryty)
+    ll_itemsize = lc.Constant.int(buf.itemsize.type, itemsize)
+    nbytes = builder.mul(buf.nitems, buf.itemsize)
+
+    # Check that the buffer size is compatible
+    rem = builder.srem(nbytes, ll_itemsize)
+    is_incompatible = cgutils.is_not_null(builder, rem)
+    with builder.if_then(is_incompatible, likely=False):
+        msg = "buffer size must be a multiple of element size"
+        context.call_conv.return_user_exc(builder, ValueError, (msg,))
+
+    shape = cgutils.pack_array(builder, [builder.sdiv(nbytes, ll_itemsize)])
+    strides = cgutils.pack_array(builder, [ll_itemsize])
+    data = builder.bitcast(buf.data,
+                           context.get_value_type(out_datamodel.get_type('data')))
+
+    populate_array(out_ary,
+                   data=data,
+                   shape=shape,
+                   strides=strides,
+                   itemsize=ll_itemsize,
+                   meminfo=buf.meminfo,
+                   parent=buf.parent,)
+
+    return out_ary._getvalue()
+
diff --git a/numba/typing/npydecl.py b/numba/typing/npydecl.py
--- a/numba/typing/npydecl.py
+++ b/numba/typing/npydecl.py
@@ -598,6 +598,28 @@ def generic(self, args, kws):
 builtin_global(numpy.linspace, types.Function(NdLinspace))
 
 
+@builtin
+class NdFromBuffer(CallableTemplate):
+    key = numpy.frombuffer
+
+    def generic(self):
+        def typer(buffer, dtype=None):
+            if not isinstance(buffer, types.Buffer) or buffer.layout != 'C':
+                return
+            if dtype is None:
+                nb_dtype = types.float64
+            else:
+                nb_dtype = _parse_dtype(dtype)
+
+            if nb_dtype is not None:
+                return types.Array(dtype=nb_dtype, ndim=1, layout='C',
+                                   readonly=not buffer.mutable)
+
+        return typer
+
+builtin_global(numpy.frombuffer, types.Function(NdFromBuffer))
+
+
 # -----------------------------------------------------------------------------
 # Miscellaneous functions
 


Test patch:
diff --git a/numba/tests/support.py b/numba/tests/support.py
--- a/numba/tests/support.py
+++ b/numba/tests/support.py
@@ -204,10 +204,15 @@ def _assertNumberEqual(first, second, delta=None):
         if compare_family == "ndarray":
             dtype = self._fix_dtype(first.dtype)
             self.assertEqual(dtype, self._fix_dtype(second.dtype))
-            self.assertEqual(first.ndim, second.ndim)
-            self.assertEqual(first.shape, second.shape)
+            self.assertEqual(first.ndim, second.ndim,
+                             "different number of dimensions")
+            self.assertEqual(first.shape, second.shape,
+                             "different shapes")
+            self.assertEqual(first.flags.writeable, second.flags.writeable,
+                             "different mutability")
             # itemsize is already checked by the dtype test above
-            self.assertEqual(self._fix_strides(first), self._fix_strides(second))
+            self.assertEqual(self._fix_strides(first), self._fix_strides(second),
+                             "different strides")
             if first.dtype != dtype:
                 first = first.astype(dtype)
             if second.dtype != dtype:
diff --git a/numba/tests/test_array_methods.py b/numba/tests/test_array_methods.py
--- a/numba/tests/test_array_methods.py
+++ b/numba/tests/test_array_methods.py
@@ -1,6 +1,7 @@
 from __future__ import division
 
 from itertools import product
+import sys
 
 import numpy as np
 
@@ -76,6 +77,26 @@ def array_view(arr):
         return arr.view(newtype)
     return array_view
 
+def np_frombuffer(b):
+    """
+    np.frombuffer() on a Python-allocated buffer.
+    """
+    return np.frombuffer(b)
+
+def np_frombuffer_dtype(b):
+    return np.frombuffer(b, dtype=np.complex64)
+
+def np_frombuffer_allocated(shape):
+    """
+    np.frombuffer() on a Numba-allocated buffer.
+    """
+    arr = np.ones(shape, dtype=np.int32)
+    return np.frombuffer(arr)
+
+def np_frombuffer_allocated_dtype(shape):
+    arr = np.ones(shape, dtype=np.int32)
+    return np.frombuffer(arr, dtype=np.complex64)
+
 
 class TestArrayMethods(TestCase):
     """
@@ -310,6 +331,63 @@ def check_err(arr, dtype):
         check_err(arr, dt1)
         check_err(arr, dt2)
 
+    @unittest.skipIf(sys.version_info < (2, 7),
+                     "buffer protocol not supported on Python 2.6")
+    def check_np_frombuffer(self, pyfunc):
+        def run(buf):
+            cres = self.ccache.compile(pyfunc, (typeof(buf),))
+            return cres.entry_point(buf)
+        def check(buf):
+            old_refcnt = sys.getrefcount(buf)
+            expected = pyfunc(buf)
+            got = run(buf)
+            self.assertPreciseEqual(got, expected)
+            del expected
+            self.assertEqual(sys.getrefcount(buf), old_refcnt + 1)
+            del got
+            self.assertEqual(sys.getrefcount(buf), old_refcnt)
+
+        b = bytearray(range(16))
+        check(b)
+        if sys.version_info >= (3,):
+            check(bytes(b))
+            check(memoryview(b))
+        check(np.arange(12))
+        b = np.arange(12).reshape((3, 4))
+        check(b)
+
+        with self.assertRaises(ValueError) as raises:
+            run(bytearray(b"xxx"))
+        self.assertEqual("buffer size must be a multiple of element size",
+                         str(raises.exception))
+
+    def test_np_frombuffer(self):
+        self.check_np_frombuffer(np_frombuffer)
+
+    def test_np_frombuffer_dtype(self):
+        self.check_np_frombuffer(np_frombuffer_dtype)
+
+    @unittest.skipIf(sys.version_info < (2, 7),
+                     "buffer protocol not supported on Python 2.6")
+    def check_np_frombuffer_allocated(self, pyfunc):
+        def run(shape):
+            cres = self.ccache.compile(pyfunc, (typeof(shape),))
+            return cres.entry_point(shape)
+        def check(shape):
+            expected = pyfunc(shape)
+            got = run(shape)
+            self.assertPreciseEqual(got, expected)
+
+        check((16,))
+        check((4, 4))
+        check((1, 0, 1))
+
+    def test_np_frombuffer_allocated(self):
+        self.check_np_frombuffer_allocated(np_frombuffer_allocated)
+
+    def test_np_frombuffer_allocated(self):
+        self.check_np_frombuffer_allocated(np_frombuffer_allocated_dtype)
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/numba/tests/test_support.py b/numba/tests/test_support.py
--- a/numba/tests/test_support.py
+++ b/numba/tests/test_support.py
@@ -243,6 +243,9 @@ def test_arrays(self):
         self.ne(a, b.T.copy().T)
         # Different ndim
         self.ne(a, b.flatten())
+        # Different writeability
+        b.flags.writeable = False
+        self.ne(a, b)
         # Precision
         a = np.arange(1, 3, dtype=np.float64)
         b = a * (1.0 + DBL_EPSILON)