Repository: numba/numba

Issue:
Keywords to DUFuncs only supported after compiled
Passing in keywords like `out` to a `DUFunc` is only supported after the first compilation. This is with numba 0.25.0.

``` python
In [37]: @nb.vectorize(nopython=True)
def foo(a, b):
    return a + b
   ....:

In [38]: a = np.array([1, 2, 3])

In [39]: o = np.zeros(3)

In [40]: foo(a, a, out=o)
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-40-abd64fdee011> in <module>()
----> 1 foo(a, a, out=o)

/Users/jcrist/anaconda/envs/dask/lib/python2.7/site-packages/numba/npyufunc/dufunc.pyc in _compile_for_args(self, *args, **kws)
    135         args_len = len(args)
    136         assert (args_len == nin) or (args_len == nin + self.ufunc.nout)
--> 137         assert not kws
    138         argtys = []
    139         # To avoid a mismatch in how Numba types values as opposed to

AssertionError:

In [41]: foo(a, a)
Out[41]: array([2, 4, 6])

In [42]: foo(a, a, out=o)
Out[42]: array([ 2.,  4.,  6.])
```



Discussion:

Resolving patch:
diff --git a/numba/npyufunc/dufunc.py b/numba/npyufunc/dufunc.py
--- a/numba/npyufunc/dufunc.py
+++ b/numba/npyufunc/dufunc.py
@@ -132,6 +132,14 @@ def add(self, sig):
 
     def _compile_for_args(self, *args, **kws):
         nin = self.ufunc.nin
+        if kws:
+            if 'out' in kws:
+                out = kws.pop('out')
+                args += (out,)
+            if kws:
+                raise TypeError("unexpected keyword arguments to ufunc: %s"
+                                % ", ".join(repr(k) for k in sorted(kws)))
+
         args_len = len(args)
         assert (args_len == nin) or (args_len == nin + self.ufunc.nout)
         assert not kws


Test patch:
diff --git a/numba/tests/npyufunc/test_ufuncbuilding.py b/numba/tests/npyufunc/test_ufuncbuilding.py
--- a/numba/tests/npyufunc/test_ufuncbuilding.py
+++ b/numba/tests/npyufunc/test_ufuncbuilding.py
@@ -9,7 +9,7 @@
 from numba import vectorize, guvectorize
 from numba.npyufunc import PyUFunc_One
 from numba.npyufunc.dufunc import DUFunc as UFuncBuilder
-from ..support import tag
+from ..support import tag, TestCase
 
 
 def add(a, b):
@@ -71,7 +71,7 @@ def guerror(a, b, c):
     raise MyException
 
 
-class TestUfuncBuilding(unittest.TestCase):
+class TestUfuncBuilding(TestCase):
 
     @tag('important')
     def test_basic_ufunc(self):
@@ -159,7 +159,7 @@ def tearDown(self):
         config.DISABLE_JIT = self.old_disable_jit
 
 
-class TestGUfuncBuilding(unittest.TestCase):
+class TestGUfuncBuilding(TestCase):
 
     def test_basic_gufunc(self):
         gufb = GUFuncBuilder(guadd, "(x, y),(x, y)->(x, y)")
@@ -214,7 +214,7 @@ def tearDown(self):
         config.DISABLE_JIT = self.old_disable_jit
 
 
-class TestVectorizeDecor(unittest.TestCase):
+class TestVectorizeDecor(TestCase):
 
     _supported_identities = [0, 1, None, "reorderable"]
 
@@ -272,6 +272,27 @@ def test_vectorize_only_kws(self):
         ufunc = vectorize(identity=PyUFunc_One, nopython=True)(mul)
         self.assertTrue(numpy.all(ufunc(a,b) == (a * b)))
 
+    def test_vectorize_output_kwarg(self):
+        """
+        Passing the output array as a keyword argument (issue #1867).
+        """
+        def check(ufunc):
+            a = numpy.arange(10, 16, dtype='int32')
+            out = numpy.zeros_like(a)
+            got = ufunc(a, a, out=out)
+            self.assertIs(got, out)
+            self.assertPreciseEqual(out, a + a)
+            with self.assertRaises(TypeError):
+                ufunc(a, a, zzz=out)
+
+        # With explicit sigs
+        ufunc = vectorize(['int32(int32, int32)'], nopython=True)(add)
+        check(ufunc)
+        # With implicit sig
+        ufunc = vectorize(nopython=True)(add)
+        check(ufunc)  # compiling
+        check(ufunc)  # after compiling
+
     @tag('important')
     def test_guvectorize(self):
         ufunc = guvectorize(['(int32[:,:], int32[:,:], int32[:,:])'],