ENH: switch ndarray.__divmod__ to use np.divmod

7 files changed, 146 insertions, 171 deletions

diff --git a/doc/release/1.13.0-notes.rst b/doc/release/1.13.0-notes.rst
index 3945639da..f7a7e0bb1 100644
--- a/doc/release/1.13.0-notes.rst
+++ b/doc/release/1.13.0-notes.rst
@@ -346,9 +346,10 @@ an error if array values do not support numeric operations.
 
 New ``divmod`` ufunc
 --------------------
-This ufunc corresponds to the Python builtin `divmod`. ``np.divmod(x, y)``
-calculates a result equivalent to
-``(np.floor_divide(x, y), np.remainder(x, y))`` but faster.
+This ufunc corresponds to the Python builtin `divmod`, and is used to implement
+`divmod` when called on numpy arrays. ``np.divmod(x, y)`` calculates a result
+equivalent to ``(np.floor_divide(x, y), np.remainder(x, y))`` but is
+approximately twice as fast as calling the functions separately.
 
 Better ``repr`` of object arrays
 --------------------------------
diff --git a/numpy/core/code_generators/ufunc_docstrings.py b/numpy/core/code_generators/ufunc_docstrings.py
index 8677e4328..b84a80a57 100644
--- a/numpy/core/code_generators/ufunc_docstrings.py
+++ b/numpy/core/code_generators/ufunc_docstrings.py
@@ -2894,7 +2894,7 @@ add_newdoc('numpy.core.umath', 'divmod',
     out1 : ndarray
         Element-wise quotient resulting from floor division.
     out2 : ndarray
-        Element-wise remainder from division.
+        Element-wise remainder from floor division.
 
     See Also
     --------
diff --git a/numpy/core/src/multiarray/number.c b/numpy/core/src/multiarray/number.c
index d6598cdb6..b8239c972 100644
--- a/numpy/core/src/multiarray/number.c
+++ b/numpy/core/src/multiarray/number.c
@@ -83,6 +83,7 @@ PyArray_SetNumericOps(PyObject *dict)
     SET(multiply);
     SET(divide);
     SET(remainder);
+    SET(divmod);
     SET(power);
     SET(square);
     SET(reciprocal);
@@ -135,6 +136,7 @@ PyArray_GetNumericOps(void)
     GET(multiply);
     GET(divide);
     GET(remainder);
+    GET(divmod);
     GET(power);
     GET(square);
     GET(reciprocal);
@@ -344,6 +346,12 @@ array_remainder(PyArrayObject *m1, PyObject *m2)
     return PyArray_GenericBinaryFunction(m1, m2, n_ops.remainder);
 }
 
+static PyObject *
+array_divmod(PyArrayObject *m1, PyObject *m2)
+{
+    BINOP_GIVE_UP_IF_NEEDED(m1, m2, nb_divmod, array_divmod);
+    return PyArray_GenericBinaryFunction(m1, m2, n_ops.divmod);
+}
 
 #if PY_VERSION_HEX >= 0x03050000
 /* Need this to be version dependent on account of the slot check */
@@ -796,37 +804,6 @@ _array_nonzero(PyArrayObject *mp)
 }
 
 
-
-static PyObject *
-array_divmod(PyArrayObject *op1, PyObject *op2)
-{
-    PyObject *divp, *modp, *result;
-
-    BINOP_GIVE_UP_IF_NEEDED(op1, op2, nb_divmod, array_divmod);
-
-    divp = array_floor_divide(op1, op2);
-    if (divp == NULL) {
-        return NULL;
-    }
-    else if(divp == Py_NotImplemented) {
-        return divp;
-    }
-    modp = array_remainder(op1, op2);
-    if (modp == NULL) {
-        Py_DECREF(divp);
-        return NULL;
-    }
-    else if(modp == Py_NotImplemented) {
-        Py_DECREF(divp);
-        return modp;
-    }
-    result = Py_BuildValue("OO", divp, modp);
-    Py_DECREF(divp);
-    Py_DECREF(modp);
-    return result;
-}
-
-
 NPY_NO_EXPORT PyObject *
 array_int(PyArrayObject *v)
 {
diff --git a/numpy/core/src/multiarray/number.h b/numpy/core/src/multiarray/number.h
index 86f681c10..113fc2475 100644
--- a/numpy/core/src/multiarray/number.h
+++ b/numpy/core/src/multiarray/number.h
@@ -7,6 +7,7 @@ typedef struct {
     PyObject *multiply;
     PyObject *divide;
     PyObject *remainder;
+    PyObject *divmod;
     PyObject *power;
     PyObject *square;
     PyObject *reciprocal;
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 3f4b183aa..089ac6d9c 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -2881,7 +2881,7 @@ class TestBinop(object):
             'truediv':  (np.true_divide, True, float),
             'floordiv': (np.floor_divide, True, float),
             'mod':      (np.remainder, True, float),
-            'divmod':   (None, False, float),
+            'divmod':   (np.divmod, False, float),
             'pow':      (np.power, True, int),
             'lshift':   (np.left_shift, True, int),
             'rshift':   (np.right_shift, True, int),
@@ -2944,13 +2944,15 @@ class TestBinop(object):
         def check(obj, binop_override_expected, ufunc_override_expected,
                   inplace_override_expected, check_scalar=True):
             for op, (ufunc, has_inplace, dtype) in ops.items():
+                err_msg = ('op: %s, ufunc: %s, has_inplace: %s, dtype: %s'
+                           % (op, ufunc, has_inplace, dtype))
                 check_objs = [np.arange(3, 5, dtype=dtype)]
                 if check_scalar:
                     check_objs.append(check_objs[0][0])
                 for arr in check_objs:
                     arr_method = getattr(arr, "__{0}__".format(op))
 
-                    def norm(result):
+                    def first_out_arg(result):
                         if op == "divmod":
                             assert_(isinstance(result, tuple))
                             return result[0]
@@ -2959,77 +2961,78 @@ class TestBinop(object):
 
                     # arr __op__ obj
                     if binop_override_expected:
-                        assert_equal(arr_method(obj), NotImplemented)
+                        assert_equal(arr_method(obj), NotImplemented, err_msg)
                     elif ufunc_override_expected:
-                        assert_equal(norm(arr_method(obj))[0],
-                                     "__array_ufunc__")
+                        assert_equal(arr_method(obj)[0], "__array_ufunc__",
+                                     err_msg)
                     else:
                         if (isinstance(obj, np.ndarray) and
                             (type(obj).__array_ufunc__ is
                              np.ndarray.__array_ufunc__)):
                             # __array__ gets ignored
-                            res = norm(arr_method(obj))
-                            assert_(res.__class__ is obj.__class__)
+                            res = first_out_arg(arr_method(obj))
+                            assert_(res.__class__ is obj.__class__, err_msg)
                         else:
                             assert_raises((TypeError, Coerced),
-                                          arr_method, obj)
+                                          arr_method, obj, err_msg=err_msg)
                     # obj __op__ arr
                     arr_rmethod = getattr(arr, "__r{0}__".format(op))
                     if ufunc_override_expected:
-                        res = norm(arr_rmethod(obj))
-                        assert_equal(res[0], "__array_ufunc__")
-                        if ufunc is not None:
-                            assert_equal(res[1], ufunc)
+                        res = arr_rmethod(obj)
+                        assert_equal(res[0], "__array_ufunc__",
+                                     err_msg=err_msg)
+                        assert_equal(res[1], ufunc, err_msg=err_msg)
                     else:
                         if (isinstance(obj, np.ndarray) and
                                 (type(obj).__array_ufunc__ is
                                  np.ndarray.__array_ufunc__)):
                             # __array__ gets ignored
-                            res = norm(arr_rmethod(obj))
-                            assert_(res.__class__ is obj.__class__)
+                            res = first_out_arg(arr_rmethod(obj))
+                            assert_(res.__class__ is obj.__class__, err_msg)
                         else:
                             # __array_ufunc__ = "asdf" creates a TypeError
                             assert_raises((TypeError, Coerced),
-                                          arr_rmethod, obj)
+                                          arr_rmethod, obj, err_msg=err_msg)
 
                     # arr __iop__ obj
                     # array scalars don't have in-place operators
                     if has_inplace and isinstance(arr, np.ndarray):
                         arr_imethod = getattr(arr, "__i{0}__".format(op))
                         if inplace_override_expected:
-                            assert_equal(arr_method(obj), NotImplemented)
+                            assert_equal(arr_method(obj), NotImplemented,
+                                         err_msg=err_msg)
                         elif ufunc_override_expected:
                             res = arr_imethod(obj)
-                            assert_equal(res[0], "__array_ufunc__")
-                            if ufunc is not None:
-                                assert_equal(res[1], ufunc)
-                                assert_(type(res[-1]["out"]) is tuple)
-                                assert_(res[-1]["out"][0] is arr)
+                            assert_equal(res[0], "__array_ufunc__", err_msg)
+                            assert_equal(res[1], ufunc, err_msg)
+                            assert_(type(res[-1]["out"]) is tuple, err_msg)
+                            assert_(res[-1]["out"][0] is arr, err_msg)
                         else:
                             if (isinstance(obj, np.ndarray) and
                                     (type(obj).__array_ufunc__ is
                                     np.ndarray.__array_ufunc__)):
                                 # __array__ gets ignored
-                                assert_(arr_imethod(obj) is arr)
+                                assert_(arr_imethod(obj) is arr, err_msg)
                             else:
                                 assert_raises((TypeError, Coerced),
-                                              arr_imethod, obj)
+                                              arr_imethod, obj,
+                                              err_msg=err_msg)
 
                     op_fn = getattr(operator, op, None)
                     if op_fn is None:
                         op_fn = getattr(operator, op + "_", None)
                     if op_fn is None:
                         op_fn = getattr(builtins, op)
-                    assert_equal(op_fn(obj, arr), "forward")
+                    assert_equal(op_fn(obj, arr), "forward", err_msg)
                     if not isinstance(obj, np.ndarray):
                         if binop_override_expected:
-                            assert_equal(op_fn(arr, obj), "reverse")
+                            assert_equal(op_fn(arr, obj), "reverse", err_msg)
                         elif ufunc_override_expected:
-                            assert_equal(norm(op_fn(arr, obj))[0],
-                                         "__array_ufunc__")
-                    if ufunc_override_expected and ufunc is not None:
-                        assert_equal(norm(ufunc(obj, arr))[0],
-                                     "__array_ufunc__")
+                            assert_equal(op_fn(arr, obj)[0], "__array_ufunc__",
+                                         err_msg)
+                    if ufunc_override_expected:
+                        assert_equal(ufunc(obj, arr)[0], "__array_ufunc__",
+                                     err_msg)
 
         # No array priority, no array_ufunc -> nothing called
         check(make_obj(object), False, False, False)
diff --git a/numpy/core/tests/test_scalarmath.py b/numpy/core/tests/test_scalarmath.py
index 1cafde5a0..c76db98f8 100644
--- a/numpy/core/tests/test_scalarmath.py
+++ b/numpy/core/tests/test_scalarmath.py
@@ -189,30 +189,34 @@ class TestPower(TestCase):
             assert_raises(TypeError, operator.pow, np.array(t(a)), b, c)
 
 
-class TestModulus(TestCase):
+def floordiv_and_mod(x, y):
+    return (x // y, x % y)
+
+
+def _signs(dt):
+    if dt in np.typecodes['UnsignedInteger']:
+        return (+1,)
+    else:
+        return (+1, -1)
 
-    floordiv = operator.floordiv
-    mod = operator.mod
+
+class TestModulus(TestCase):
 
     def test_modulus_basic(self):
         dt = np.typecodes['AllInteger'] + np.typecodes['Float']
-        for dt1, dt2 in itertools.product(dt, dt):
-            for sg1, sg2 in itertools.product((+1, -1), (+1, -1)):
-                if sg1 == -1 and dt1 in np.typecodes['UnsignedInteger']:
-                    continue
-                if sg2 == -1 and dt2 in np.typecodes['UnsignedInteger']:
-                    continue
-                fmt = 'dt1: %s, dt2: %s, sg1: %s, sg2: %s'
-                msg = fmt % (dt1, dt2, sg1, sg2)
-                a = np.array(sg1*71, dtype=dt1)[()]
-                b = np.array(sg2*19, dtype=dt2)[()]
-                div = self.floordiv(a, b)
-                rem = self.mod(a, b)
-                assert_equal(div*b + rem, a, err_msg=msg)
-                if sg2 == -1:
-                    assert_(b < rem <= 0, msg)
-                else:
-                    assert_(b > rem >= 0, msg)
+        for op in [floordiv_and_mod, divmod]:
+            for dt1, dt2 in itertools.product(dt, dt):
+                for sg1, sg2 in itertools.product(_signs(dt1), _signs(dt2)):
+                    fmt = 'op: %s, dt1: %s, dt2: %s, sg1: %s, sg2: %s'
+                    msg = fmt % (op.__name__, dt1, dt2, sg1, sg2)
+                    a = np.array(sg1*71, dtype=dt1)[()]
+                    b = np.array(sg2*19, dtype=dt2)[()]
+                    div, rem = op(a, b)
+                    assert_equal(div*b + rem, a, err_msg=msg)
+                    if sg2 == -1:
+                        assert_(b < rem <= 0, msg)
+                    else:
+                        assert_(b > rem >= 0, msg)
 
     def test_float_modulus_exact(self):
         # test that float results are exact for small integers. This also
@@ -231,42 +235,42 @@ class TestModulus(TestCase):
         tgtdiv = np.where((tgtdiv == 0.0) & ((b < 0) ^ (a < 0)), -0.0, tgtdiv)
         tgtrem = np.where((tgtrem == 0.0) & (b < 0), -0.0, tgtrem)
 
-        for dt in np.typecodes['Float']:
-            msg = 'dtype: %s' % (dt,)
-            fa = a.astype(dt)
-            fb = b.astype(dt)
-            # use list comprehension so a_ and b_ are scalars
-            div = [self.floordiv(a_, b_) for  a_, b_ in zip(fa, fb)]
-            rem = [self.mod(a_, b_) for a_, b_ in zip(fa, fb)]
-            assert_equal(div, tgtdiv, err_msg=msg)
-            assert_equal(rem, tgtrem, err_msg=msg)
+        for op in [floordiv_and_mod, divmod]:
+            for dt in np.typecodes['Float']:
+                msg = 'op: %s, dtype: %s' % (op.__name__, dt)
+                fa = a.astype(dt)
+                fb = b.astype(dt)
+                # use list comprehension so a_ and b_ are scalars
+                div, rem = zip(*[op(a_, b_) for  a_, b_ in zip(fa, fb)])
+                assert_equal(div, tgtdiv, err_msg=msg)
+                assert_equal(rem, tgtrem, err_msg=msg)
 
     def test_float_modulus_roundoff(self):
         # gh-6127
         dt = np.typecodes['Float']
-        for dt1, dt2 in itertools.product(dt, dt):
-            for sg1, sg2 in itertools.product((+1, -1), (+1, -1)):
-                fmt = 'dt1: %s, dt2: %s, sg1: %s, sg2: %s'
-                msg = fmt % (dt1, dt2, sg1, sg2)
-                a = np.array(sg1*78*6e-8, dtype=dt1)[()]
-                b = np.array(sg2*6e-8, dtype=dt2)[()]
-                div = self.floordiv(a, b)
-                rem = self.mod(a, b)
-                # Equal assertion should hold when fmod is used
-                assert_equal(div*b + rem, a, err_msg=msg)
-                if sg2 == -1:
-                    assert_(b < rem <= 0, msg)
-                else:
-                    assert_(b > rem >= 0, msg)
+        for op in [floordiv_and_mod, divmod]:
+            for dt1, dt2 in itertools.product(dt, dt):
+                for sg1, sg2 in itertools.product((+1, -1), (+1, -1)):
+                    fmt = 'op: %s, dt1: %s, dt2: %s, sg1: %s, sg2: %s'
+                    msg = fmt % (op.__name__, dt1, dt2, sg1, sg2)
+                    a = np.array(sg1*78*6e-8, dtype=dt1)[()]
+                    b = np.array(sg2*6e-8, dtype=dt2)[()]
+                    div, rem = op(a, b)
+                    # Equal assertion should hold when fmod is used
+                    assert_equal(div*b + rem, a, err_msg=msg)
+                    if sg2 == -1:
+                        assert_(b < rem <= 0, msg)
+                    else:
+                        assert_(b > rem >= 0, msg)
 
     def test_float_modulus_corner_cases(self):
         # Check remainder magnitude.
         for dt in np.typecodes['Float']:
             b = np.array(1.0, dtype=dt)
             a = np.nextafter(np.array(0.0, dtype=dt), -b)
-            rem = self.mod(a, b)
+            rem = operator.mod(a, b)
             assert_(rem <= b, 'dt: %s' % dt)
-            rem = self.mod(-a, -b)
+            rem = operator.mod(-a, -b)
             assert_(rem >= -b, 'dt: %s' % dt)
 
         # Check nans, inf
@@ -277,14 +281,14 @@ class TestModulus(TestCase):
                 fzer = np.array(0.0, dtype=dt)
                 finf = np.array(np.inf, dtype=dt)
                 fnan = np.array(np.nan, dtype=dt)
-                rem = self.mod(fone, fzer)
+                rem = operator.mod(fone, fzer)
                 assert_(np.isnan(rem), 'dt: %s' % dt)
                 # MSVC 2008 returns NaN here, so disable the check.
-                #rem = self.mod(fone, finf)
+                #rem = operator.mod(fone, finf)
                 #assert_(rem == fone, 'dt: %s' % dt)
-                rem = self.mod(fone, fnan)
+                rem = operator.mod(fone, fnan)
                 assert_(np.isnan(rem), 'dt: %s' % dt)
-                rem = self.mod(finf, fone)
+                rem = operator.mod(finf, fone)
                 assert_(np.isnan(rem), 'dt: %s' % dt)
 
 
diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py
index 7d147f727..51bf7c942 100644
--- a/numpy/core/tests/test_umath.py
+++ b/numpy/core/tests/test_umath.py
@@ -264,34 +264,34 @@ class TestDivision(TestCase):
         assert_equal(y, [1.e+110, 0], err_msg=msg)
 
 
+def floor_divide_and_remainder(x, y):
+    return (np.floor_divide(x, y), np.remainder(x, y))
+
+
+def _signs(dt):
+    if dt in np.typecodes['UnsignedInteger']:
+        return (+1,)
+    else:
+        return (+1, -1)
+
+
 class TestRemainder(TestCase):
 
     def test_remainder_basic(self):
         dt = np.typecodes['AllInteger'] + np.typecodes['Float']
-        for dt1, dt2 in itertools.product(dt, dt):
-            for sg1, sg2 in itertools.product((+1, -1), (+1, -1)):
-                if sg1 == -1 and dt1 in np.typecodes['UnsignedInteger']:
-                    continue
-                if sg2 == -1 and dt2 in np.typecodes['UnsignedInteger']:
-                    continue
-                fmt = 'dt1: %s, dt2: %s, sg1: %s, sg2: %s'
-                msg = fmt % (dt1, dt2, sg1, sg2)
-                a = np.array(sg1*71, dtype=dt1)
-                b = np.array(sg2*19, dtype=dt2)
-                div = np.floor_divide(a, b)
-                rem = np.remainder(a, b)
-                assert_equal(div*b + rem, a, err_msg=msg)
-                if sg2 == -1:
-                    assert_(b < rem <= 0, msg)
-                else:
-                    assert_(b > rem >= 0, msg)
-
-                div, rem = np.divmod(a, b)
-                assert_equal(div*b + rem, a, err_msg=msg)
-                if sg2 == -1:
-                    assert_(b < rem <= 0, msg)
-                else:
-                    assert_(b > rem >= 0, msg)
+        for op in [floor_divide_and_remainder, np.divmod]:
+            for dt1, dt2 in itertools.product(dt, dt):
+                for sg1, sg2 in itertools.product(_signs(dt1), _signs(dt2)):
+                    fmt = 'op: %s, dt1: %s, dt2: %s, sg1: %s, sg2: %s'
+                    msg = fmt % (op.__name__, dt1, dt2, sg1, sg2)
+                    a = np.array(sg1*71, dtype=dt1)
+                    b = np.array(sg2*19, dtype=dt2)
+                    div, rem = op(a, b)
+                    assert_equal(div*b + rem, a, err_msg=msg)
+                    if sg2 == -1:
+                        assert_(b < rem <= 0, msg)
+                    else:
+                        assert_(b > rem >= 0, msg)
 
     def test_float_remainder_exact(self):
         # test that float results are exact for small integers. This also
@@ -310,43 +310,32 @@ class TestRemainder(TestCase):
         tgtdiv = np.where((tgtdiv == 0.0) & ((b < 0) ^ (a < 0)), -0.0, tgtdiv)
         tgtrem = np.where((tgtrem == 0.0) & (b < 0), -0.0, tgtrem)
 
-        for dt in np.typecodes['Float']:
-            msg = 'dtype: %s' % (dt,)
-            fa = a.astype(dt)
-            fb = b.astype(dt)
-            div = np.floor_divide(fa, fb)
-            rem = np.remainder(fa, fb)
-            assert_equal(div, tgtdiv, err_msg=msg)
-            assert_equal(rem, tgtrem, err_msg=msg)
-
-            div, rem = np.divmod(fa, fb)
-            assert_equal(div, tgtdiv, err_msg=msg)
-            assert_equal(rem, tgtrem, err_msg=msg)
+        for op in [floor_divide_and_remainder, np.divmod]:
+            for dt in np.typecodes['Float']:
+                msg = 'op: %s, dtype: %s' % (op.__name__, dt)
+                fa = a.astype(dt)
+                fb = b.astype(dt)
+                div, rem = op(fa, fb)
+                assert_equal(div, tgtdiv, err_msg=msg)
+                assert_equal(rem, tgtrem, err_msg=msg)
 
     def test_float_remainder_roundoff(self):
         # gh-6127
         dt = np.typecodes['Float']
-        for dt1, dt2 in itertools.product(dt, dt):
-            for sg1, sg2 in itertools.product((+1, -1), (+1, -1)):
-                fmt = 'dt1: %s, dt2: %s, sg1: %s, sg2: %s'
-                msg = fmt % (dt1, dt2, sg1, sg2)
-                a = np.array(sg1*78*6e-8, dtype=dt1)
-                b = np.array(sg2*6e-8, dtype=dt2)
-                div = np.floor_divide(a, b)
-                rem = np.remainder(a, b)
-                # Equal assertion should hold when fmod is used
-                assert_equal(div*b + rem, a, err_msg=msg)
-                if sg2 == -1:
-                    assert_(b < rem <= 0, msg)
-                else:
-                    assert_(b > rem >= 0, msg)
-
-                div, rem = np.divmod(a, b)
-                assert_equal(div*b + rem, a, err_msg=msg)
-                if sg2 == -1:
-                    assert_(b < rem <= 0, msg)
-                else:
-                    assert_(b > rem >= 0, msg)
+        for op in [floor_divide_and_remainder, np.divmod]:
+            for dt1, dt2 in itertools.product(dt, dt):
+                for sg1, sg2 in itertools.product((+1, -1), (+1, -1)):
+                    fmt = 'op: %s, dt1: %s, dt2: %s, sg1: %s, sg2: %s'
+                    msg = fmt % (op.__name__, dt1, dt2, sg1, sg2)
+                    a = np.array(sg1*78*6e-8, dtype=dt1)
+                    b = np.array(sg2*6e-8, dtype=dt2)
+                    div, rem = op(a, b)
+                    # Equal assertion should hold when fmod is used
+                    assert_equal(div*b + rem, a, err_msg=msg)
+                    if sg2 == -1:
+                        assert_(b < rem <= 0, msg)
+                    else:
+                        assert_(b > rem >= 0, msg)
 
     def test_float_remainder_corner_cases(self):
         # Check remainder magnitude.