Skip to content

implement maybe_wrap_in_index for EA Mixin dispatch #21715

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.

Sign up for GitHub

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

Closed
wants to merge 3 commits into from
Closed

implement maybe_wrap_in_index for EA Mixin dispatch #21715

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
725e239
implement maybe_wrap_in_index
jbrockmendel Jul 3, 2018
1b01b3a
remove unused imports
jbrockmendel Jul 3, 2018
4c6f25e
Merge branch 'master' of https://github.com/pandas-dev/pandas into dt…
jbrockmendel Jul 3, 2018
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
16 changes: 15 additions & 1 deletion pandas/core/arrays/datetimes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,10 @@
import warnings

import numpy as np
from pytz import utc

from pandas._libs.tslib import Timestamp, NaT, iNaT
from pandas._libs.tslibs import timezones
from pandas._libs.tslibs import timezones, conversion

from pandas.util._decorators import cache_readonly

Expand Down Expand Up @@ -108,3 +109,16 @@ def _sub_datelike_dti(self, other):
mask = (self._isnan) | (other._isnan)
new_values[mask] = iNaT
return new_values.view('timedelta64[ns]')

# -----------------------------------------------------------------
# Timezone Conversion and Localization Methods

def _local_timestamps(self):
values = self.asi8
indexer = values.argsort()
result = conversion.tz_convert(values.take(indexer), utc, self.tz)

n = len(indexer)
reverse = np.empty(n, dtype=np.int_)
reverse.put(indexer, np.arange(n))
return result.take(reverse)
36 changes: 29 additions & 7 deletions pandas/core/indexes/datetimelike.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -314,13 +314,9 @@ def _evaluate_compare(self, other, op):
mask = (self._isnan) | (other._isnan)
if is_bool_dtype(result):
result[mask] = False
return result

result[mask] = iNaT
try:
return Index(result)
except TypeError:
return result
else:
result[mask] = iNaT
return maybe_wrap_in_index(result, strict=False)

def _ensure_localized(self, result):
"""
Expand Down Expand Up @@ -1137,3 +1133,29 @@ def _ensure_datetimelike_to_i8(other):
# period array cannot be coerces to int
other = Index(other).asi8
return other


def maybe_wrap_in_index(result, name=None, strict=True):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so this doesn't belong here, prob should just be in .base itself. why do you need the strict parameter? that is really odd

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

indexes.base makes sense, sure.

The strict parameter is there because most of the use cases don't have the try/except for the index wrapping. strict=False is only passed for the comparison ops. I'll see if I can come up with something prettier. At least this has served its purpose of tagging the next few things to be moved.

"""
After dispatching to a DatetimelikeArrayMixin method we get an ndarray
back. If possible, wrap this array in an Index. The Index constructor
will find the appropriate Index subclass.

Parameters
----------
result : ndarray
name : object, default None
strict : bool, default True

Returns
-------
out : ndarray or Index
"""
if is_bool_dtype(result):
return result
try:
return Index(result, name=name)
except TypeError:
if not strict:
return result
raise
24 changes: 7 additions & 17 deletions pandas/core/indexes/datetimes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
is_integer_dtype,
is_datetime64_ns_dtype, is_datetimelike,
is_period_dtype,
is_bool_dtype,
is_string_like,
is_list_like,
is_scalar,
Expand All @@ -42,7 +41,7 @@
import pandas.compat as compat
from pandas.tseries.frequencies import to_offset, get_period_alias, Resolution
from pandas.core.indexes.datetimelike import (
DatelikeOps, TimelikeOps, DatetimeIndexOpsMixin)
DatelikeOps, TimelikeOps, DatetimeIndexOpsMixin, maybe_wrap_in_index)
from pandas.tseries.offsets import (
DateOffset, generate_range, Tick, CDay, prefix_mapping)

Expand Down Expand Up @@ -95,7 +94,7 @@ def f(self):
result = fields.get_date_field(values, field)
result = self._maybe_mask_results(result, convert='float64')

return Index(result, name=self.name)
return maybe_wrap_in_index(result, name=self.name)

f.__name__ = name
f.__doc__ = docstring
Expand Down Expand Up @@ -149,10 +148,7 @@ def wrapper(self, other):
if self.hasnans:
result[self._isnan] = nat_result

# support of bool dtype indexers
if is_bool_dtype(result):
return result
return Index(result)
return maybe_wrap_in_index(result)

return compat.set_function_name(wrapper, opname, cls)

Expand Down Expand Up @@ -603,14 +599,8 @@ def _local_timestamps(self):
if self.is_monotonic:
return conversion.tz_convert(self.asi8, utc, self.tz)
else:
values = self.asi8
indexer = values.argsort()
result = conversion.tz_convert(values.take(indexer), utc, self.tz)

n = len(indexer)
reverse = np.empty(n, dtype=np.int_)
reverse.put(indexer, np.arange(n))
return result.take(reverse)
# fall back to non-optimized implementation
return DatetimeArrayMixin._local_timestamps(self)

@classmethod
def _simple_new(cls, values, name=None, freq=None, tz=None,
Expand Down Expand Up @@ -2438,7 +2428,7 @@ def month_name(self, locale=None):
result = fields.get_date_name_field(values, 'month_name',
locale=locale)
result = self._maybe_mask_results(result)
return Index(result, name=self.name)
return maybe_wrap_in_index(result, name=self.name)

def day_name(self, locale=None):
"""
Expand All @@ -2464,7 +2454,7 @@ def day_name(self, locale=None):
result = fields.get_date_name_field(values, 'day_name',
locale=locale)
result = self._maybe_mask_results(result)
return Index(result, name=self.name)
return maybe_wrap_in_index(result, name=self.name)


DatetimeIndex._add_comparison_methods()
Expand Down
8 changes: 5 additions & 3 deletions pandas/core/indexes/period.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,8 @@
from pandas.tseries.offsets import Tick, DateOffset

from pandas.core.indexes.datetimes import DatetimeIndex, Int64Index, Index
from pandas.core.indexes.datetimelike import DatelikeOps, DatetimeIndexOpsMixin
from pandas.core.indexes.datetimelike import (
DatelikeOps, DatetimeIndexOpsMixin, maybe_wrap_in_index)
from pandas.core.tools.datetimes import parse_time_string

from pandas._libs.lib import infer_dtype
Expand Down Expand Up @@ -55,7 +56,8 @@ def _field_accessor(name, alias, docstring=None):
def f(self):
base, mult = _gfc(self.freq)
result = get_period_field_arr(alias, self._ndarray_values, base)
return Index(result, name=self.name)
return maybe_wrap_in_index(result, name=self.name)

f.__name__ = name
f.__doc__ = docstring
return property(f)
Expand Down Expand Up @@ -761,7 +763,7 @@ def _sub_period(self, other):
if self.hasnans:
new_data[self._isnan] = tslib.NaT

return Index(new_data)
return maybe_wrap_in_index(new_data)

def shift(self, n):
"""
Expand Down
16 changes: 6 additions & 10 deletions pandas/core/indexes/timedeltas.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
_TD_DTYPE,
is_integer,
is_float,
is_bool_dtype,
is_list_like,
is_scalar,
is_timedelta64_dtype,
Expand All @@ -30,7 +29,7 @@
import pandas.core.dtypes.concat as _concat
from pandas.util._decorators import Appender, Substitution, deprecate_kwarg
from pandas.core.indexes.datetimelike import (
TimelikeOps, DatetimeIndexOpsMixin)
TimelikeOps, DatetimeIndexOpsMixin, maybe_wrap_in_index)
from pandas.core.tools.timedeltas import (
to_timedelta, _coerce_scalar_to_timedelta_type)
from pandas.tseries.offsets import Tick, DateOffset
Expand All @@ -47,7 +46,7 @@ def f(self):
if self.hasnans:
result = self._maybe_mask_results(result, convert='float64')

return Index(result, name=self.name)
return maybe_wrap_in_index(result, name=self.name)

f.__name__ = name
f.__doc__ = docstring
Expand Down Expand Up @@ -89,10 +88,7 @@ def wrapper(self, other):
if self.hasnans:
result[self._isnan] = nat_result

# support of bool dtype indexers
if is_bool_dtype(result):
return result
return Index(result)
return maybe_wrap_in_index(result)

return compat.set_function_name(wrapper, opname, cls)

Expand Down Expand Up @@ -409,7 +405,7 @@ def _evaluate_with_timedelta_like(self, other, op):
else:
result = op(left, np.float64(right))
result = self._maybe_mask_results(result, convert='float64')
return Index(result, name=self.name, copy=False)
return maybe_wrap_in_index(result, name=self.name)

return NotImplemented

Expand Down Expand Up @@ -549,8 +545,8 @@ def total_seconds(self):
Float64Index([0.0, 86400.0, 172800.0, 259200.00000000003, 345600.0],
dtype='float64')
"""
return Index(self._maybe_mask_results(1e-9 * self.asi8),
name=self.name)
result = self._maybe_mask_results(1e-9 * self.asi8)
return maybe_wrap_in_index(result, name=self.name)

def to_pytimedelta(self):
"""
Expand Down