MAINT: lib: A few changes in some vectorize tests.

Added a note to some of the tests that the order of the calls is an
important part of the test, because the code is testing the caching
of the ufuncs.

Also reuse a couple functions in the vectorize tests by defining them at
the module level.

1 files changed, 31 insertions, 41 deletions

diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index 9470a355d..b4e928273 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -3,6 +3,7 @@ import warnings
 import sys
 import decimal
 from fractions import Fraction
+import math
 import pytest
 
 import numpy as np
@@ -1221,6 +1222,16 @@ class TestExtins:
         assert_array_equal(a, ac)
 
 
+# _foo1 and _foo2 are used in some tests in TestVectorize.
+
+def _foo1(x, y=1.0):
+    return y*math.floor(x)
+
+
+def _foo2(x, y=1.0, z=0.0):
+    return y*math.floor(x) + z
+
+
 class TestVectorize:
 
     def test_simple(self):
@@ -1252,7 +1263,6 @@ class TestVectorize:
         assert_array_equal(y, x)
 
     def test_ufunc(self):
-        import math
         f = vectorize(math.cos)
         args = np.array([0, 0.5 * np.pi, np.pi, 1.5 * np.pi, 2 * np.pi])
         r1 = f(args)
@@ -1276,13 +1286,9 @@ class TestVectorize:
     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])
+        f = vectorize(_foo1, otypes=[float])
+        # We're testing the caching of ufuncs by vectorize, so the order
+        # of these function calls is an important part of the test.
         r1 = f(np.arange(3.0), 1.0)
         r2 = f(np.arange(3.0))
         assert_array_equal(r1, r2)
@@ -1291,13 +1297,9 @@ class TestVectorize:
         # 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])
+        f = vectorize(_foo1, otypes=[float])
+        # We're testing the caching of ufuncs by vectorize, so the order
+        # of these function calls is an important part of the test.
         r1 = f(np.arange(3.0))
         r2 = f(np.arange(3.0), 1.0)
         assert_array_equal(r1, r2)
@@ -1305,13 +1307,9 @@ class TestVectorize:
     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])
+        f = vectorize(_foo1, otypes=[float])
+        # We're testing the caching of ufuncs by vectorize, so the order
+        # of these function calls is an important part of the test.
         r1 = f(np.arange(3.0))
         r2 = f(np.arange(3.0), y=1.0)
         r3 = f(np.arange(3.0))
@@ -1321,34 +1319,26 @@ class TestVectorize:
     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])
+        f = vectorize(_foo2, otypes=[float])
+        # We're testing the caching of ufuncs by vectorize, so the order
+        # of these function calls is an important part of the test.
         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))
+        assert_equal(r1, _foo2(10.4, z=100))
+        assert_equal(r2, _foo2(10.4, y=-1))
+        assert_equal(r3, _foo2(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])
+        f = vectorize(_foo2, otypes=[float])
+        # We're testing the caching of ufuncs by vectorize, so the order
+        # of these function calls is an important part of the test.
         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))
+        assert_equal(r1, _foo2(z=100, x=10.4, y=-1))
+        assert_equal(r2, _foo2(1, 2, 3))
 
     def test_keywords_no_func_code(self):
         # This needs to test a function that has keywords but