Skip to content

Commit

Permalink
TYP: define RangeIndex methods non-dynamically (#36931)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel committed Oct 8, 2020
1 parent c04e231 commit 91e8bad
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 72 deletions.
149 changes: 80 additions & 69 deletions pandas/core/indexes/range.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -812,83 +811,95 @@ def any(self, *args, **kwargs) -> bool:

# --------------------------------------------------------------------

@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)
3 changes: 0 additions & 3 deletions setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -193,9 +193,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

Expand Down

0 comments on commit 91e8bad

Please sign in to comment.