ENH: Add divmod to series and index. (#14208)

Author: llllllllll

doc/source/basics.rst

@@ -188,6 +188,32 @@ And similarly for ``axis="items"`` and ``axis="minor"``.
    match the broadcasting behavior of Panel. Though it would require a
    transition period so users can change their code...
 
+Series and Index also support the :func:`divmod` builtin. This function takes
+the floor division and modulo operation at the same time returning a two-tuple
+of the same type as the left hand side. For example:
+
+.. ipython:: python
+
+   s = pd.Series(np.arange(10))
+   s
+   div, rem = divmod(s, 3)
+   div
+   rem
+
+   idx = pd.Index(np.arange(10))
+   idx
+   div, rem = divmod(idx, 3)
+   div
+   rem
+
+We can also do elementwise :func:`divmod`:
+
+.. ipython:: python
+
+   div, rem = divmod(s, [2, 2, 3, 3, 4, 4, 5, 5, 6, 6])
+   div
+   rem
+
 Missing data / operations with fill values
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

doc/source/whatsnew/v0.19.0.txt

@@ -1330,6 +1330,9 @@ Other API Changes
 - ``pd.read_csv()`` in the C engine will now issue a ``ParserWarning`` or raise a ``ValueError`` when ``sep`` encoded is more than one character long (:issue:`14065`)
 - ``DataFrame.values`` will now return ``float64`` with a ``DataFrame`` of mixed ``int64`` and ``uint64`` dtypes, conforming to ``np.find_common_type`` (:issue:`10364`, :issue:`13917`)
 - ``pd.read_stata()`` can now handle some format 111 files, which are produced by SAS when generating Stata dta files (:issue:`11526`)
+- ``Series`` and ``Index`` now support ``divmod`` which will return a tuple of
+  series or indices. This behaves like a standard binary operator with regards
+  to broadcasting rules (:issue:`14208`).
 
 .. _whatsnew_0190.deprecations:
 

pandas/core/ops.py

@@ -39,7 +39,8 @@
 
 
 def _create_methods(arith_method, comp_method, bool_method,
-                    use_numexpr, special=False, default_axis='columns'):
+                    use_numexpr, special=False, default_axis='columns',
+                    have_divmod=False):
     # creates actual methods based upon arithmetic, comp and bool method
     # constructors.
 
@@ -127,6 +128,15 @@ def names(x):
                                   names('ror_'), op('|')),
                  rxor=bool_method(lambda x, y: operator.xor(y, x),
                                   names('rxor'), op('^'))))
+    if have_divmod:
+        # divmod doesn't have an op that is supported by numexpr
+        new_methods['divmod'] = arith_method(
+            divmod,
+            names('divmod'),
+            None,
+            default_axis=default_axis,
+            construct_result=_construct_divmod_result,
+        )
 
     new_methods = dict((names(k), v) for k, v in new_methods.items())
     return new_methods
@@ -156,7 +166,7 @@ def add_methods(cls, new_methods, force, select, exclude):
 def add_special_arithmetic_methods(cls, arith_method=None,
                                    comp_method=None, bool_method=None,
                                    use_numexpr=True, force=False, select=None,
-                                   exclude=None):
+                                   exclude=None, have_divmod=False):
     """
     Adds the full suite of special arithmetic methods (``__add__``,
     ``__sub__``, etc.) to the class.
@@ -177,14 +187,17 @@ def add_special_arithmetic_methods(cls, arith_method=None,
         if passed, only sets functions with names in select
     exclude : iterable of strings (optional)
         if passed, will not set functions with names in exclude
+    have_divmod : bool, (optional)
+        should a divmod method be added? this method is special because it
+        returns a tuple of cls instead of a single element of type cls
     """
 
     # in frame, special methods have default_axis = None, comp methods use
     # 'columns'
 
     new_methods = _create_methods(arith_method, comp_method,
                                   bool_method, use_numexpr, default_axis=None,
-                                  special=True)
+                                  special=True, have_divmod=have_divmod)
 
     # inplace operators (I feel like these should get passed an `inplace=True`
     # or just be removed
@@ -618,8 +631,22 @@ def _align_method_SERIES(left, right, align_asobject=False):
     return left, right
 
 
+def _construct_result(left, result, index, name, dtype):
+    return left._constructor(result, index=index, name=name, dtype=dtype)
+
+
+def _construct_divmod_result(left, result, index, name, dtype):
+    """divmod returns a tuple of like indexed series instead of a single series.
+    """
+    constructor = left._constructor
+    return (
+        constructor(result[0], index=index, name=name, dtype=dtype),
+        constructor(result[1], index=index, name=name, dtype=dtype),
+    )
+
+
 def _arith_method_SERIES(op, name, str_rep, fill_zeros=None, default_axis=None,
-                         **eval_kwargs):
+                         construct_result=_construct_result, **eval_kwargs):
     """
     Wrapper function for Series arithmetic operations, to avoid
     code duplication.
@@ -692,8 +719,14 @@ def wrapper(left, right, name=name, na_op=na_op):
                 lvalues = lvalues.values
 
         result = wrap_results(safe_na_op(lvalues, rvalues))
-        return left._constructor(result, index=left.index,
-                                 name=name, dtype=dtype)
+        return construct_result(
+            left,
+            result,
+            index=left.index,
+            name=name,
+            dtype=dtype,
+        )
+
     return wrapper
 
 
@@ -933,6 +966,10 @@ def wrapper(self, other):
                                  'desc': 'Integer division',
                                  'reversed': False,
                                  'reverse': 'rfloordiv'},
+                    'divmod': {'op': 'divmod',
+                               'desc': 'Integer division and modulo',
+                               'reversed': False,
+                               'reverse': None},
 
                     'eq': {'op': '==',
                                  'desc': 'Equal to',
@@ -1033,7 +1070,8 @@ def flex_wrapper(self, other, level=None, fill_value=None, axis=0):
 
 series_special_funcs = dict(arith_method=_arith_method_SERIES,
                             comp_method=_comp_method_SERIES,
-                            bool_method=_bool_method_SERIES)
+                            bool_method=_bool_method_SERIES,
+                            have_divmod=True)
 
 _arith_doc_FRAME = """
 Binary operator %s with support to substitute a fill_value for missing data in

pandas/indexes/base.py

@@ -3426,7 +3426,7 @@ def _validate_for_numeric_binop(self, other, op, opstr):
     def _add_numeric_methods_binary(cls):
         """ add in numeric methods """
 
-        def _make_evaluate_binop(op, opstr, reversed=False):
+        def _make_evaluate_binop(op, opstr, reversed=False, constructor=Index):
             def _evaluate_numeric_binop(self, other):
 
                 from pandas.tseries.offsets import DateOffset
@@ -3448,7 +3448,7 @@ def _evaluate_numeric_binop(self, other):
                 attrs = self._maybe_update_attributes(attrs)
                 with np.errstate(all='ignore'):
                     result = op(values, other)
-                return Index(result, **attrs)
+                return constructor(result, **attrs)
 
             return _evaluate_numeric_binop
 
@@ -3478,6 +3478,15 @@ def _evaluate_numeric_binop(self, other):
             cls.__rdiv__ = _make_evaluate_binop(
                 operator.div, '__div__', reversed=True)
 
+        cls.__divmod__ = _make_evaluate_binop(
+            divmod,
+            '__divmod__',
+            constructor=lambda result, **attrs: (
+                Index(result[0], **attrs),
+                Index(result[1], **attrs),
+            ),
+        )
+
     @classmethod
     def _add_numeric_methods_unary(cls):
         """ add in numeric unary methods """

pandas/tests/indexes/test_numeric.py

@@ -73,6 +73,30 @@ def test_numeric_compat(self):
         self.assertRaises(ValueError, lambda: idx * idx[0:3])
         self.assertRaises(ValueError, lambda: idx * np.array([1, 2]))
 
+        result = divmod(idx, 2)
+        with np.errstate(all='ignore'):
+            div, mod = divmod(idx.values, 2)
+            expected = Index(div), Index(mod)
+        for r, e in zip(result, expected):
+            tm.assert_index_equal(r, e)
+
+        result = divmod(idx, np.full_like(idx.values, 2))
+        with np.errstate(all='ignore'):
+            div, mod = divmod(idx.values, np.full_like(idx.values, 2))
+            expected = Index(div), Index(mod)
+        for r, e in zip(result, expected):
+            tm.assert_index_equal(r, e)
+
+        result = divmod(idx, Series(np.full_like(idx.values, 2)))
+        with np.errstate(all='ignore'):
+            div, mod = divmod(
+                idx.values,
+                np.full_like(idx.values, 2),
+            )
+            expected = Index(div), Index(mod)
+        for r, e in zip(result, expected):
+            tm.assert_index_equal(r, e)
+
     def test_explicit_conversions(self):
 
         # GH 8608

pandas/tests/series/test_operators.py

@@ -1,6 +1,7 @@
 # coding=utf-8
 # pylint: disable-msg=E1101,W0612
 
+from collections import Iterable
 from datetime import datetime, timedelta
 import operator
 from itertools import product, starmap
@@ -19,7 +20,7 @@
 from pandas.compat import range, zip
 from pandas import compat
 from pandas.util.testing import (assert_series_equal, assert_almost_equal,
-                                 assert_frame_equal)
+                                 assert_frame_equal, assert_index_equal)
 import pandas.util.testing as tm
 
 from .common import TestData
@@ -185,6 +186,34 @@ def check_comparators(series, other, check_dtype=True):
         check_comparators(self.ts, 5)
         check_comparators(self.ts, self.ts + 1, check_dtype=False)
 
+    def test_divmod(self):
+        def check(series, other):
+            results = divmod(series, other)
+            if isinstance(other, Iterable) and len(series) != len(other):
+                # if the lengths don't match, this is the test where we use
+                # `self.ts[::2]`. Pad every other value in `other_np` with nan.
+                other_np = []
+                for n in other:
+                    other_np.append(n)
+                    other_np.append(np.nan)
+            else:
+                other_np = other
+            other_np = np.asarray(other_np)
+            with np.errstate(all='ignore'):
+                expecteds = divmod(series.values, np.asarray(other_np))
+
+            for result, expected in zip(results, expecteds):
+                # check the values, name, and index separatly
+                assert_almost_equal(np.asarray(result), expected)
+
+                self.assertEqual(result.name, series.name)
+                assert_index_equal(result.index, series.index)
+
+        check(self.ts, self.ts * 2)
+        check(self.ts, self.ts * 0)
+        check(self.ts, self.ts[::2])
+        check(self.ts, 5)
+
     def test_operators_empty_int_corner(self):
         s1 = Series([], [], dtype=np.int32)
         s2 = Series({'x': 0.})