From a6d5c0f550fb64cf495b3874c1d9b875a7181744 Mon Sep 17 00:00:00 2001 From: Brock Date: Tue, 6 Oct 2020 17:45:41 -0700 Subject: [PATCH] TYP: define RangeIndex methods non-dynamically --- pandas/core/indexes/range.py | 149 +++++++++++++++++++---------------- setup.cfg | 3 - 2 files changed, 80 insertions(+), 72 deletions(-) diff --git a/pandas/core/indexes/range.py b/pandas/core/indexes/range.py index d5d9a9b5bc0a3..8d6d26a2d01e4 100644 --- a/pandas/core/indexes/range.py +++ b/pandas/core/indexes/range.py @@ -9,7 +9,6 @@ from pandas._libs import index as libindex from pandas._libs.lib import no_default from pandas._typing import Label -import pandas.compat as compat from pandas.compat.numpy import function as nv from pandas.util._decorators import Appender, cache_readonly, doc @@ -752,83 +751,95 @@ def all(self) -> bool: def any(self) -> bool: return any(self._range) - @classmethod - def _add_numeric_methods_binary(cls): - """ add in numeric methods, specialized to RangeIndex """ - - def _make_evaluate_binop(op, step=False): - """ - Parameters - ---------- - op : callable that accepts 2 params - perform the binary op - step : callable, optional, default to False - op to apply to the step parm if not None - if False, use the existing step - """ - - @unpack_zerodim_and_defer(op.__name__) - def _evaluate_numeric_binop(self, other): - if isinstance(other, ABCTimedeltaIndex): - # Defer to TimedeltaIndex implementation - return NotImplemented - elif isinstance(other, (timedelta, np.timedelta64)): - # GH#19333 is_integer evaluated True on timedelta64, - # so we need to catch these explicitly - return op(self._int64index, other) - elif is_timedelta64_dtype(other): - # Must be an np.ndarray; GH#22390 - return op(self._int64index, other) - - other = extract_array(other, extract_numpy=True) - attrs = self._get_attributes_dict() - - left, right = self, other + def _arith_method(self, other, op, step=False): + """ + Parameters + ---------- + other : Any + op : callable that accepts 2 params + perform the binary op + step : callable, optional, default to False + op to apply to the step parm if not None + if False, use the existing step + """ + + if isinstance(other, ABCTimedeltaIndex): + # Defer to TimedeltaIndex implementation + return NotImplemented + elif isinstance(other, (timedelta, np.timedelta64)): + # GH#19333 is_integer evaluated True on timedelta64, + # so we need to catch these explicitly + return op(self._int64index, other) + elif is_timedelta64_dtype(other): + # Must be an np.ndarray; GH#22390 + return op(self._int64index, other) + + other = extract_array(other, extract_numpy=True) + attrs = self._get_attributes_dict() + + left, right = self, other - try: - # apply if we have an override - if step: - with np.errstate(all="ignore"): - rstep = step(left.step, right) + try: + # apply if we have an override + if step: + with np.errstate(all="ignore"): + rstep = step(left.step, right) + + # we don't have a representable op + # so return a base index + if not is_integer(rstep) or not rstep: + raise ValueError + + else: + rstep = left.step + + with np.errstate(all="ignore"): + rstart = op(left.start, right) + rstop = op(left.stop, right) + + result = type(self)(rstart, rstop, rstep, **attrs) - # we don't have a representable op - # so return a base index - if not is_integer(rstep) or not rstep: - raise ValueError + # for compat with numpy / Int64Index + # even if we can represent as a RangeIndex, return + # as a Float64Index if we have float-like descriptors + if not all(is_integer(x) for x in [rstart, rstop, rstep]): + result = result.astype("float64") - else: - rstep = left.step + return result - with np.errstate(all="ignore"): - rstart = op(left.start, right) - rstop = op(left.stop, right) + except (ValueError, TypeError, ZeroDivisionError): + # Defer to Int64Index implementation + return op(self._int64index, other) + # TODO: Do attrs get handled reliably? - result = type(self)(rstart, rstop, rstep, **attrs) + @unpack_zerodim_and_defer("__add__") + def __add__(self, other): + return self._arith_method(other, operator.add) - # for compat with numpy / Int64Index - # even if we can represent as a RangeIndex, return - # as a Float64Index if we have float-like descriptors - if not all(is_integer(x) for x in [rstart, rstop, rstep]): - result = result.astype("float64") + @unpack_zerodim_and_defer("__radd__") + def __radd__(self, other): + return self._arith_method(other, ops.radd) - return result + @unpack_zerodim_and_defer("__sub__") + def __sub__(self, other): + return self._arith_method(other, operator.sub) - except (ValueError, TypeError, ZeroDivisionError): - # Defer to Int64Index implementation - return op(self._int64index, other) - # TODO: Do attrs get handled reliably? + @unpack_zerodim_and_defer("__rsub__") + def __rsub__(self, other): + return self._arith_method(other, ops.rsub) - name = f"__{op.__name__}__" - return compat.set_function_name(_evaluate_numeric_binop, name, cls) + @unpack_zerodim_and_defer("__mul__") + def __mul__(self, other): + return self._arith_method(other, operator.mul, step=operator.mul) - cls.__add__ = _make_evaluate_binop(operator.add) - cls.__radd__ = _make_evaluate_binop(ops.radd) - cls.__sub__ = _make_evaluate_binop(operator.sub) - cls.__rsub__ = _make_evaluate_binop(ops.rsub) - cls.__mul__ = _make_evaluate_binop(operator.mul, step=operator.mul) - cls.__rmul__ = _make_evaluate_binop(ops.rmul, step=ops.rmul) - cls.__truediv__ = _make_evaluate_binop(operator.truediv, step=operator.truediv) - cls.__rtruediv__ = _make_evaluate_binop(ops.rtruediv, step=ops.rtruediv) + @unpack_zerodim_and_defer("__rmul__") + def __rmul__(self, other): + return self._arith_method(other, ops.rmul, step=ops.rmul) + @unpack_zerodim_and_defer("__truediv__") + def __truediv__(self, other): + return self._arith_method(other, operator.truediv, step=operator.truediv) -RangeIndex._add_numeric_methods() + @unpack_zerodim_and_defer("__rtruediv__") + def __rtruediv__(self, other): + return self._arith_method(other, ops.rtruediv, step=ops.rtruediv) diff --git a/setup.cfg b/setup.cfg index 4c6c7f6d9598c..25dcff93b0e12 100644 --- a/setup.cfg +++ b/setup.cfg @@ -196,9 +196,6 @@ check_untyped_defs=False [mypy-pandas.core.indexes.multi] check_untyped_defs=False -[mypy-pandas.core.indexes.range] -check_untyped_defs=False - [mypy-pandas.core.internals.blocks] check_untyped_defs=False