From 08b39390a0984748dcc08bfd83c3f45b7b43f554 Mon Sep 17 00:00:00 2001 From: Joe Jevnik Date: Mon, 12 Sep 2016 14:53:15 -0400 Subject: [PATCH] ENH: Add divmod to series and index. --- doc/source/basics.rst | 26 ++++++++++++++ doc/source/whatsnew/v0.19.0.txt | 3 ++ pandas/core/ops.py | 52 +++++++++++++++++++++++---- pandas/indexes/base.py | 13 +++++-- pandas/tests/indexes/test_numeric.py | 24 +++++++++++++ pandas/tests/series/test_operators.py | 31 +++++++++++++++- 6 files changed, 139 insertions(+), 10 deletions(-) diff --git a/doc/source/basics.rst b/doc/source/basics.rst index 1f670fb7fb593..19318aad3d53d 100644 --- a/doc/source/basics.rst +++ b/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 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/doc/source/whatsnew/v0.19.0.txt b/doc/source/whatsnew/v0.19.0.txt index f3a6736ff9920..ffb6e72019602 100644 --- a/doc/source/whatsnew/v0.19.0.txt +++ b/doc/source/whatsnew/v0.19.0.txt @@ -1328,6 +1328,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: diff --git a/pandas/core/ops.py b/pandas/core/ops.py index b81d62c3cda18..237b9394dfc25 100644 --- a/pandas/core/ops.py +++ b/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,6 +187,9 @@ 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 @@ -184,7 +197,7 @@ def add_special_arithmetic_methods(cls, arith_method=None, 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 diff --git a/pandas/indexes/base.py b/pandas/indexes/base.py index d4ca18a6713b5..f430305f5cb91 100644 --- a/pandas/indexes/base.py +++ b/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 """ diff --git a/pandas/tests/indexes/test_numeric.py b/pandas/tests/indexes/test_numeric.py index d3a89b301ae46..51d8c95f9d783 100644 --- a/pandas/tests/indexes/test_numeric.py +++ b/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 diff --git a/pandas/tests/series/test_operators.py b/pandas/tests/series/test_operators.py index 197311868b768..24c26276ea24d 100644 --- a/pandas/tests/series/test_operators.py +++ b/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.})