Permit non intp typed ints as args in np.* allocators.

numba/targets/arrayobj.py

@@ -3149,7 +3149,7 @@ def iternext_numpy_ndindex(context, builder, sig, args, result):
 # -----------------------------------------------------------------------------
 # Numpy array constructors
 
-def _empty_nd_impl(context, builder, arrtype, shapes):
+def _empty_nd_impl(context, builder, arrtype, shapes, shapety=None):
     """Utility function used for allocating a new array during LLVM code
     generation (lowering).  Given a target context, builder, array
     type, and a tuple or list of lowered dimension sizes, returns a
@@ -3162,20 +3162,38 @@ def _empty_nd_impl(context, builder, arrtype, shapes):
     itemsize = context.get_constant(types.intp, get_itemsize(context, arrtype))
 
     # compute array length
+    maxbits = types.intp.bitwidth
+    cast_shapes = []
+    if shapety is None:
+        [cast_shapes.append(x) for x in shapes]
+    else: # typed shapes
+        if isinstance(shapety, types.BaseTuple):
+            for i in range(shapety.count):
+                ty = shapety[i]
+                if ty.bitwidth > maxbits:
+                    msg = ("Type %s given in shape would lead to allocation "
+                           "potentially exceeding intp size" % ty)
+                    raise ValueError(msg)
+                val = shapes[i]
+                cast_shapes.append(context.cast(builder, val, ty, types.intp))
+        else:
+            cast_shapes.append(context.cast(builder, shapes[0], shapety,
+                                            types.intp))
+
     arrlen = context.get_constant(types.intp, 1)
-    for s in shapes:
+    for s in cast_shapes:
         arrlen = builder.mul(arrlen, s)
 
     if arrtype.ndim == 0:
         strides = ()
     elif arrtype.layout == 'C':
         strides = [itemsize]
-        for dimension_size in reversed(shapes[1:]):
+        for dimension_size in reversed(cast_shapes[1:]):
             strides.append(builder.mul(strides[-1], dimension_size))
         strides = tuple(reversed(strides))
     elif arrtype.layout == 'F':
         strides = [itemsize]
-        for dimension_size in shapes[:-1]:
+        for dimension_size in cast_shapes[:-1]:
             strides.append(builder.mul(strides[-1], dimension_size))
         strides = tuple(strides)
     else:
@@ -3191,7 +3209,7 @@ def _empty_nd_impl(context, builder, arrtype, shapes):
     data = context.nrt.meminfo_data(builder, meminfo)
 
     intp_t = context.get_value_type(types.intp)
-    shape_array = cgutils.pack_array(builder, shapes, ty=intp_t)
+    shape_array = cgutils.pack_array(builder, cast_shapes, ty=intp_t)
     strides_array = cgutils.pack_array(builder, strides, ty=intp_t)
 
     populate_array(ary,
@@ -3241,7 +3259,8 @@ def _parse_empty_args(context, builder, sig, args):
     arrshapetype = sig.args[0]
     arrshape = args[0]
     arrtype = sig.return_type
-    return arrtype, _parse_shape(context, builder, arrshapetype, arrshape)
+    return arrtype, _parse_shape(context, builder, arrshapetype, arrshape), \
+           arrshapetype
 
 
 def _parse_empty_like_args(context, builder, sig, args):
@@ -3261,8 +3280,8 @@ def _parse_empty_like_args(context, builder, sig, args):
 @lower_builtin(np.empty, types.Any)
 @lower_builtin(np.empty, types.Any, types.Any)
 def numpy_empty_nd(context, builder, sig, args):
-    arrtype, shapes = _parse_empty_args(context, builder, sig, args)
-    ary = _empty_nd_impl(context, builder, arrtype, shapes)
+    arrtype, shapes, shapety = _parse_empty_args(context, builder, sig, args)
+    ary = _empty_nd_impl(context, builder, arrtype, shapes, shapety)
     return impl_ret_new_ref(context, builder, sig.return_type, ary._getvalue())
 
 @lower_builtin(np.empty_like, types.Any)
@@ -3276,8 +3295,8 @@ def numpy_empty_like_nd(context, builder, sig, args):
 @lower_builtin(np.zeros, types.Any)
 @lower_builtin(np.zeros, types.Any, types.Any)
 def numpy_zeros_nd(context, builder, sig, args):
-    arrtype, shapes = _parse_empty_args(context, builder, sig, args)
-    ary = _empty_nd_impl(context, builder, arrtype, shapes)
+    arrtype, shapes, shapety = _parse_empty_args(context, builder, sig, args)
+    ary = _empty_nd_impl(context, builder, arrtype, shapes, shapety)
     _zero_fill_array(context, builder, ary)
     return impl_ret_new_ref(context, builder, sig.return_type, ary._getvalue())
 

numba/tests/test_dyn_array.py

@@ -7,13 +7,13 @@
 import threading
 
 from numba import unittest_support as unittest
-from numba.errors import TypingError
+from numba.errors import TypingError, LoweringError
 from numba import njit
 from numba import utils
 from numba.numpy_support import version as numpy_version
+from numba.config import IS_32BITS
 from .support import MemoryLeakMixin, TestCase, tag
 
-
 nrtjit = njit(_nrt=True, nogil=True)
 
 
@@ -104,6 +104,27 @@ def foo(n):
         np.testing.assert_equal(123, arr)
         del arr
 
+    def test_empty_1d_typed_arg(self):
+        @nrtjit
+        def foo(n):
+            arr = np.empty(np.int8(n))
+            for i in range(n):
+                arr[i] = i
+
+            return arr
+
+        n = 3
+        arr = foo(n)
+        self.assert_array_nrt_refct(arr, 1)
+        np.testing.assert_equal(np.arange(n), arr)
+        self.assertEqual(arr.size, n)
+        self.assertEqual(arr.shape, (n,))
+        self.assertEqual(arr.dtype, np.dtype(np.float64))
+        self.assertEqual(arr.strides, (np.dtype(np.float64).itemsize,))
+        arr.fill(123)  # test writability
+        np.testing.assert_equal(123, arr)
+        del arr
+
     def test_empty_2d(self):
         def pyfunc(m, n):
             arr = np.empty((m, n), np.int32)
@@ -153,6 +174,55 @@ def pyfunc(m, n, p):
 
         del got_arr
 
+    @tag('important')
+    def test_empty2d_typed_tuple_arg(self):
+        def pyfunc(m, n):
+            arr = np.empty((m, n), np.int32)
+            for i in range(m):
+                for j in range(n):
+                    arr[i, j] = i + j
+
+            return arr
+
+        cfunc = nrtjit(pyfunc)
+        m = np.int8(4)
+        n = np.int16(3)
+        expected_arr = pyfunc(m, n)
+        got_arr = cfunc(m, n)
+        self.assert_array_nrt_refct(got_arr, 1)
+        np.testing.assert_equal(expected_arr, got_arr)
+
+        self.assertEqual(expected_arr.size, got_arr.size)
+        self.assertEqual(expected_arr.shape, got_arr.shape)
+        self.assertEqual(expected_arr.strides, got_arr.strides)
+
+        del got_arr
+
+    @tag('important')
+    @unittest.skipUnless(IS_32BITS, "32 bit only test")
+    def test_empty2d_typed_tuple_arg_overflow(self):
+        def pyfunc(m, n):
+            arr = np.empty((m, n), np.int32)
+            for i in range(m):
+                for j in range(n):
+                    arr[i, j] = i + j
+
+            return arr
+
+        cfunc = nrtjit(pyfunc)
+        m = np.int64(4)
+        n = np.int64(3)
+
+        self.disable_leak_check()
+
+        msg = ("Type int64 given in shape would lead to allocation "
+                "potentially exceeding intp size")
+
+        with self.assertRaises(LoweringError) as raises:
+            cfunc(m, n)
+
+        self.assertIn(msg, str(raises.exception))
+
     @tag('important')
     def test_empty_2d_sliced(self):
         def pyfunc(m, n, p):
@@ -594,6 +664,12 @@ def func(n):
             return pyfunc(n)
         self.check_1d(func)
 
+    def test_1d_typed_value(self):
+        pyfunc = self.pyfunc
+        def func(n):
+            return pyfunc(np.int8(n))
+        self.check_1d(func)
+
     def test_1d_dtype(self):
         pyfunc = self.pyfunc
         def func(n):
@@ -614,6 +690,12 @@ def func(m, n):
             return pyfunc((m, n))
         self.check_2d(func)
 
+    def test_2d_typed_tuple(self):
+        pyfunc = self.pyfunc
+        def func(m, n):
+            return pyfunc((np.int8(m), np.int16(n)))
+        self.check_2d(func)
+
     @tag('important')
     def test_2d_dtype_kwarg(self):
         pyfunc = self.pyfunc