|
1 | 1 | # -*- coding: utf-8 -*-
|
| 2 | +from __future__ import division |
| 3 | + |
2 | 4 | from datetime import timedelta
|
| 5 | +import operator |
3 | 6 | import warnings
|
4 | 7 |
|
5 | 8 | import numpy as np
|
|
16 | 19 | _TD_DTYPE, ensure_int64, is_datetime64_dtype, is_float_dtype,
|
17 | 20 | is_integer_dtype, is_list_like, is_object_dtype, is_string_dtype,
|
18 | 21 | is_timedelta64_dtype)
|
19 |
| -from pandas.core.dtypes.generic import ABCSeries, ABCTimedeltaIndex |
| 22 | +from pandas.core.dtypes.generic import ( |
| 23 | + ABCDataFrame, ABCIndexClass, ABCSeries, ABCTimedeltaIndex) |
20 | 24 | from pandas.core.dtypes.missing import isna
|
21 | 25 |
|
| 26 | +from pandas.core import ops |
22 | 27 | from pandas.core.algorithms import checked_add_with_arr
|
23 | 28 | import pandas.core.common as com
|
24 | 29 |
|
@@ -101,8 +106,32 @@ def wrapper(self, other):
|
101 | 106 | return compat.set_function_name(wrapper, opname, cls)
|
102 | 107 |
|
103 | 108 |
|
| 109 | +def _wrap_tdi_op(op): |
| 110 | + """ |
| 111 | + Instead of re-implementing multiplication/division etc operations |
| 112 | + in the Array class, for now we dispatch to the TimedeltaIndex |
| 113 | + implementations. |
| 114 | + """ |
| 115 | + # TODO: implement directly here and wrap in TimedeltaIndex, instead of |
| 116 | + # the other way around |
| 117 | + def method(self, other): |
| 118 | + if isinstance(other, (ABCSeries, ABCDataFrame, ABCIndexClass)): |
| 119 | + return NotImplemented |
| 120 | + |
| 121 | + from pandas import TimedeltaIndex |
| 122 | + obj = TimedeltaIndex(self) |
| 123 | + result = op(obj, other) |
| 124 | + if is_timedelta64_dtype(result): |
| 125 | + return type(self)(result) |
| 126 | + return np.array(result) |
| 127 | + |
| 128 | + method.__name__ = '__{name}__'.format(name=op.__name__) |
| 129 | + return method |
| 130 | + |
| 131 | + |
104 | 132 | class TimedeltaArrayMixin(dtl.DatetimeLikeArrayMixin):
|
105 | 133 | _typ = "timedeltaarray"
|
| 134 | + __array_priority__ = 1000 |
106 | 135 |
|
107 | 136 | @property
|
108 | 137 | def _box_func(self):
|
@@ -293,11 +322,25 @@ def _evaluate_with_timedelta_like(self, other, op):
|
293 | 322 |
|
294 | 323 | return NotImplemented
|
295 | 324 |
|
| 325 | + __mul__ = _wrap_tdi_op(operator.mul) |
| 326 | + __rmul__ = __mul__ |
| 327 | + __truediv__ = _wrap_tdi_op(operator.truediv) |
| 328 | + __floordiv__ = _wrap_tdi_op(operator.floordiv) |
| 329 | + __rfloordiv__ = _wrap_tdi_op(ops.rfloordiv) |
| 330 | + |
| 331 | + if compat.PY2: |
| 332 | + __div__ = __truediv__ |
| 333 | + |
| 334 | + # Note: TimedeltaIndex overrides this in call to cls._add_numeric_methods |
296 | 335 | def __neg__(self):
|
297 | 336 | if self.freq is not None:
|
298 | 337 | return type(self)(-self._data, freq=-self.freq)
|
299 | 338 | return type(self)(-self._data)
|
300 | 339 |
|
| 340 | + def __abs__(self): |
| 341 | + # Note: freq is not preserved |
| 342 | + return type(self)(np.abs(self._data)) |
| 343 | + |
301 | 344 | # ----------------------------------------------------------------
|
302 | 345 | # Conversion Methods - Vectorized analogues of Timedelta methods
|
303 | 346 |
|
|
0 commit comments