From 684e4a92d5f8a9b8744ea7f994d636e7483e530f Mon Sep 17 00:00:00 2001
From: Warren Weckesser <warren.weckesser@gmail.com>
Date: Thu, 30 Apr 2020 22:29:18 -0400
Subject: BUG: lib: Fix a problem with vectorize with default parameters.

When `otypes` is given to `vectorize` and then the instance is
called, it creates a ufunc by calling numpy.core.umath.frompyfunc.
The number of arguments given to this ufunc is set to the number
of arguments in the call of the vectorize instance.  This ufunc
is cached, so frompyfunc does not have to be called on the next
call. The problem is that, if the function being wrapped has
parameters with default values, the number of arguments passed
to the vectorize instance can change, and when that happens, a
new ufunc must be created by calling frompyfunc with the correct
number of arguments.

This commit changes the cache of the ufunc from a simple attribute
that holds the most recent ufunc to a dictionary whose keys are
the number of arguments in the call.  The cache is only used when
the vectorized function is called with only positional arguments
and there are no excluded arguments.  If keywords are used, the
number of arguments is no longer sufficient to uniquely identify a
previously created ufunc.

Closes gh-16120.
---
 numpy/lib/tests/test_function_base.py | 77 +++++++++++++++++++++++++++++++++++
 1 file changed, 77 insertions(+)

(limited to 'numpy/lib/tests/test_function_base.py')

diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index 23bf3296d..9470a355d 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -1273,6 +1273,83 @@ class TestVectorize:
         r2 = np.array([3, 4, 5])
         assert_array_equal(r1, r2)
 
+    def test_keywords_with_otypes_order1(self):
+        # gh-1620: The second call of f would crash with
+        # `ValueError: invalid number of arguments`.
+
+        import math
+
+        def foo(x, y=1.0):
+            return y*math.floor(x)
+
+        f = vectorize(foo, otypes=[float])
+        r1 = f(np.arange(3.0), 1.0)
+        r2 = f(np.arange(3.0))
+        assert_array_equal(r1, r2)
+
+    def test_keywords_with_otypes_order2(self):
+        # gh-1620: The second call of f would crash with
+        # `ValueError: non-broadcastable output operand with shape ()
+        # doesn't match the broadcast shape (3,)`.
+
+        import math
+
+        def foo(x, y=1.0):
+            return y*math.floor(x)
+
+        f = vectorize(foo, otypes=[float])
+        r1 = f(np.arange(3.0))
+        r2 = f(np.arange(3.0), 1.0)
+        assert_array_equal(r1, r2)
+
+    def test_keywords_with_otypes_order3(self):
+        # gh-1620: The third call of f would crash with
+        # `ValueError: invalid number of arguments`.
+
+        import math
+
+        def foo(x, y=1.0):
+            return y*math.floor(x)
+
+        f = vectorize(foo, otypes=[float])
+        r1 = f(np.arange(3.0))
+        r2 = f(np.arange(3.0), y=1.0)
+        r3 = f(np.arange(3.0))
+        assert_array_equal(r1, r2)
+        assert_array_equal(r1, r3)
+
+    def test_keywords_with_otypes_several_kwd_args1(self):
+        # gh-1620 Make sure different uses of keyword arguments
+        # don't break the vectorized function.
+
+        import math
+
+        def foo(x, y=1.0, z=0.0):
+            return y*math.floor(x) + z
+
+        f = vectorize(foo, otypes=[float])
+        r1 = f(10.4, z=100)
+        r2 = f(10.4, y=-1)
+        r3 = f(10.4)
+        assert_equal(r1, foo(10.4, z=100))
+        assert_equal(r2, foo(10.4, y=-1))
+        assert_equal(r3, foo(10.4))
+
+    def test_keywords_with_otypes_several_kwd_args2(self):
+        # gh-1620 Make sure different uses of keyword arguments
+        # don't break the vectorized function.
+
+        import math
+
+        def foo(x, y=1.0, z=0.0):
+            return y*math.floor(x) + z
+
+        f = vectorize(foo, otypes=[float])
+        r1 = f(z=100, x=10.4, y=-1)
+        r2 = f(1, 2, 3)
+        assert_equal(r1, foo(z=100, x=10.4, y=-1))
+        assert_equal(r2, foo(1, 2, 3))
+
     def test_keywords_no_func_code(self):
         # This needs to test a function that has keywords but
         # no func_code attribute, since otherwise vectorize will
-- 
cgit v1.2.1