Skip to content

REF: share TimedeltaArray division code #50441

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Dec 28, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
230 changes: 97 additions & 133 deletions pandas/core/arrays/timedeltas.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

from datetime import timedelta
import operator
from typing import (
TYPE_CHECKING,
Iterator,
Expand Down Expand Up @@ -65,7 +66,7 @@
from pandas.core.arrays import datetimelike as dtl
from pandas.core.arrays._ranges import generate_regular_range
import pandas.core.common as com
from pandas.core.construction import extract_array
from pandas.core.ops import roperator
from pandas.core.ops.common import unpack_zerodim_and_defer

if TYPE_CHECKING:
Expand Down Expand Up @@ -492,10 +493,11 @@ def __mul__(self, other) -> TimedeltaArray:

__rmul__ = __mul__

@unpack_zerodim_and_defer("__truediv__")
def __truediv__(self, other):
# timedelta / X is well-defined for timedelta-like or numeric X

def _scalar_divlike_op(self, other, op):
"""
Shared logic for __truediv__, __rtruediv__, __floordiv__, __rfloordiv__
with scalar 'other'.
"""
if isinstance(other, self._recognized_scalars):
other = Timedelta(other)
# mypy assumes that __new__ returns an instance of the class
Expand All @@ -507,31 +509,86 @@ def __truediv__(self, other):
return result

# otherwise, dispatch to Timedelta implementation
return self._ndarray / other
return op(self._ndarray, other)

elif lib.is_scalar(other):
# assume it is numeric
result = self._ndarray / other
else:
# caller is responsible for checking lib.is_scalar(other)
# assume other is numeric, otherwise numpy will raise

if op in [roperator.rtruediv, roperator.rfloordiv]:
raise TypeError(
f"Cannot divide {type(other).__name__} by {type(self).__name__}"
)

result = op(self._ndarray, other)
freq = None

if self.freq is not None:
# Tick division is not implemented, so operate on Timedelta
freq = self.freq.delta / other
freq = to_offset(freq)
# Note: freq gets division, not floor-division, even if op
# is floordiv.
freq = self.freq / other

# TODO: 2022-12-24 test_ufunc_coercions, test_tdi_ops_attributes
# get here for truediv, no tests for floordiv

if op is operator.floordiv:
if freq.nanos == 0 and self.freq.nanos != 0:
# e.g. if self.freq is Nano(1) then dividing by 2
# rounds down to zero
# TODO: 2022-12-24 should implement the same check
# for truediv case
freq = None

return type(self)._simple_new(result, dtype=result.dtype, freq=freq)

def _cast_divlike_op(self, other):
if not hasattr(other, "dtype"):
# e.g. list, tuple
other = np.array(other)

if len(other) != len(self):
raise ValueError("Cannot divide vectors with unequal lengths")
return other

if is_timedelta64_dtype(other.dtype):
# let numpy handle it
return self._ndarray / other
def _vector_divlike_op(self, other, op) -> np.ndarray | TimedeltaArray:
"""
Shared logic for __truediv__, __floordiv__, and their reversed versions
with timedelta64-dtype ndarray other.
"""
# Let numpy handle it
result = op(self._ndarray, np.asarray(other))

elif is_object_dtype(other.dtype):
other = extract_array(other, extract_numpy=True)
if (is_integer_dtype(other.dtype) or is_float_dtype(other.dtype)) and op in [
operator.truediv,
operator.floordiv,
]:
return type(self)._simple_new(result, dtype=result.dtype)

if op in [operator.floordiv, roperator.rfloordiv]:
mask = self.isna() | isna(other)
if mask.any():
result = result.astype(np.float64)
np.putmask(result, mask, np.nan)

return result

@unpack_zerodim_and_defer("__truediv__")
def __truediv__(self, other):
# timedelta / X is well-defined for timedelta-like or numeric X
op = operator.truediv
if is_scalar(other):
return self._scalar_divlike_op(other, op)

other = self._cast_divlike_op(other)
if (
is_timedelta64_dtype(other.dtype)
or is_integer_dtype(other.dtype)
or is_float_dtype(other.dtype)
):
return self._vector_divlike_op(other, op)

if is_object_dtype(other.dtype):
other = np.asarray(other)
if self.ndim > 1:
res_cols = [left / right for left, right in zip(self, other)]
res_cols2 = [x.reshape(1, -1) for x in res_cols]
Expand All @@ -542,40 +599,18 @@ def __truediv__(self, other):
return result

else:
result = self._ndarray / other
return type(self)._simple_new(result, dtype=result.dtype)
return NotImplemented

@unpack_zerodim_and_defer("__rtruediv__")
def __rtruediv__(self, other):
# X / timedelta is defined only for timedelta-like X
if isinstance(other, self._recognized_scalars):
other = Timedelta(other)
# mypy assumes that __new__ returns an instance of the class
# github.com/python/mypy/issues/1020
if cast("Timedelta | NaTType", other) is NaT:
# specifically timedelta64-NaT
result = np.empty(self.shape, dtype=np.float64)
result.fill(np.nan)
return result

# otherwise, dispatch to Timedelta implementation
return other / self._ndarray

elif lib.is_scalar(other):
raise TypeError(
f"Cannot divide {type(other).__name__} by {type(self).__name__}"
)

if not hasattr(other, "dtype"):
# e.g. list, tuple
other = np.array(other)

if len(other) != len(self):
raise ValueError("Cannot divide vectors with unequal lengths")
op = roperator.rtruediv
if is_scalar(other):
return self._scalar_divlike_op(other, op)

other = self._cast_divlike_op(other)
if is_timedelta64_dtype(other.dtype):
# let numpy handle it
return other / self._ndarray
return self._vector_divlike_op(other, op)

elif is_object_dtype(other.dtype):
# Note: unlike in __truediv__, we do not _need_ to do type
Expand All @@ -585,60 +620,24 @@ def __rtruediv__(self, other):
return np.array(result_list)

else:
raise TypeError(
f"Cannot divide {other.dtype} data by {type(self).__name__}"
)
return NotImplemented

@unpack_zerodim_and_defer("__floordiv__")
def __floordiv__(self, other):

op = operator.floordiv
if is_scalar(other):
if isinstance(other, self._recognized_scalars):
other = Timedelta(other)
# mypy assumes that __new__ returns an instance of the class
# github.com/python/mypy/issues/1020
if cast("Timedelta | NaTType", other) is NaT:
# treat this specifically as timedelta-NaT
result = np.empty(self.shape, dtype=np.float64)
result.fill(np.nan)
return result

# dispatch to Timedelta implementation
return other.__rfloordiv__(self._ndarray)

# at this point we should only have numeric scalars; anything
# else will raise
result = self._ndarray // other
freq = None
if self.freq is not None:
# Note: freq gets division, not floor-division
freq = self.freq / other
if freq.nanos == 0 and self.freq.nanos != 0:
# e.g. if self.freq is Nano(1) then dividing by 2
# rounds down to zero
freq = None
return type(self)(result, freq=freq)

if not hasattr(other, "dtype"):
# list, tuple
other = np.array(other)
if len(other) != len(self):
raise ValueError("Cannot divide with unequal lengths")
return self._scalar_divlike_op(other, op)

if is_timedelta64_dtype(other.dtype):
other = type(self)(other)

# numpy timedelta64 does not natively support floordiv, so operate
# on the i8 values
result = self.asi8 // other.asi8
mask = self._isnan | other._isnan
if mask.any():
result = result.astype(np.float64)
np.putmask(result, mask, np.nan)
return result
other = self._cast_divlike_op(other)
if (
is_timedelta64_dtype(other.dtype)
or is_integer_dtype(other.dtype)
or is_float_dtype(other.dtype)
):
return self._vector_divlike_op(other, op)

elif is_object_dtype(other.dtype):
other = extract_array(other, extract_numpy=True)
other = np.asarray(other)
if self.ndim > 1:
res_cols = [left // right for left, right in zip(self, other)]
res_cols2 = [x.reshape(1, -1) for x in res_cols]
Expand All @@ -649,61 +648,26 @@ def __floordiv__(self, other):
assert result.dtype == object
return result

elif is_integer_dtype(other.dtype) or is_float_dtype(other.dtype):
result = self._ndarray // other
return type(self)(result)

else:
dtype = getattr(other, "dtype", type(other).__name__)
raise TypeError(f"Cannot divide {dtype} by {type(self).__name__}")
return NotImplemented

@unpack_zerodim_and_defer("__rfloordiv__")
def __rfloordiv__(self, other):

op = roperator.rfloordiv
if is_scalar(other):
if isinstance(other, self._recognized_scalars):
other = Timedelta(other)
# mypy assumes that __new__ returns an instance of the class
# github.com/python/mypy/issues/1020
if cast("Timedelta | NaTType", other) is NaT:
# treat this specifically as timedelta-NaT
result = np.empty(self.shape, dtype=np.float64)
result.fill(np.nan)
return result

# dispatch to Timedelta implementation
return other.__floordiv__(self._ndarray)

raise TypeError(
f"Cannot divide {type(other).__name__} by {type(self).__name__}"
)

if not hasattr(other, "dtype"):
# list, tuple
other = np.array(other)

if len(other) != len(self):
raise ValueError("Cannot divide with unequal lengths")
return self._scalar_divlike_op(other, op)

other = self._cast_divlike_op(other)
if is_timedelta64_dtype(other.dtype):
other = type(self)(other)
# numpy timedelta64 does not natively support floordiv, so operate
# on the i8 values
result = other.asi8 // self.asi8
mask = self._isnan | other._isnan
if mask.any():
result = result.astype(np.float64)
np.putmask(result, mask, np.nan)
return result
return self._vector_divlike_op(other, op)

elif is_object_dtype(other.dtype):
result_list = [other[n] // self[n] for n in range(len(self))]
result = np.array(result_list)
return result

else:
dtype = getattr(other, "dtype", type(other).__name__)
raise TypeError(f"Cannot divide {dtype} by {type(self).__name__}")
return NotImplemented

@unpack_zerodim_and_defer("__mod__")
def __mod__(self, other):
Expand Down
4 changes: 3 additions & 1 deletion pandas/tests/arithmetic/test_numeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -177,10 +177,12 @@ def test_div_td64arr(self, left, box_cls):
result = right // left
tm.assert_equal(result, expected)

msg = "Cannot divide"
# (true_) needed for min-versions build 2022-12-26
msg = "ufunc '(true_)?divide' cannot use operands with types"
with pytest.raises(TypeError, match=msg):
left / right

msg = "ufunc 'floor_divide' cannot use operands with types"
with pytest.raises(TypeError, match=msg):
left // right

Expand Down
1 change: 1 addition & 0 deletions pandas/tests/arithmetic/test_timedelta64.py
Original file line number Diff line number Diff line change
Expand Up @@ -2019,6 +2019,7 @@ def test_td64arr_div_numeric_array(
"cannot perform __truediv__",
"unsupported operand",
"Cannot divide",
"ufunc 'divide' cannot use operands with types",
]
)
with pytest.raises(TypeError, match=pattern):
Expand Down