|
41 | 41 | is_dtype_equal, |
42 | 42 | is_integer, |
43 | 43 | is_list_like, |
44 | | - is_period_dtype, |
45 | 44 | ) |
46 | 45 | from pandas.core.dtypes.concat import concat_compat |
47 | 46 |
|
@@ -101,23 +100,6 @@ class DatetimeIndexOpsMixin(NDArrayBackedExtensionIndex): |
101 | 100 | def _is_all_dates(self) -> bool: |
102 | 101 | return True |
103 | 102 |
|
104 | | - # ------------------------------------------------------------------------ |
105 | | - # Abstract data attributes |
106 | | - |
107 | | - @property |
108 | | - def values(self) -> np.ndarray: |
109 | | - # Note: PeriodArray overrides this to return an ndarray of objects. |
110 | | - return self._data._ndarray |
111 | | - |
112 | | - def __array_wrap__(self, result, context=None): |
113 | | - """ |
114 | | - Gets called after a ufunc and other functions. |
115 | | - """ |
116 | | - out = super().__array_wrap__(result, context=context) |
117 | | - if isinstance(out, DatetimeTimedeltaMixin) and self.freq is not None: |
118 | | - out = out._with_freq("infer") |
119 | | - return out |
120 | | - |
121 | 103 | # ------------------------------------------------------------------------ |
122 | 104 |
|
123 | 105 | def equals(self, other: Any) -> bool: |
@@ -165,21 +147,6 @@ def __contains__(self, key: Any) -> bool: |
165 | 147 | return False |
166 | 148 | return True |
167 | 149 |
|
168 | | - @Appender(_index_shared_docs["take"] % _index_doc_kwargs) |
169 | | - def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs): |
170 | | - nv.validate_take((), kwargs) |
171 | | - indices = np.asarray(indices, dtype=np.intp) |
172 | | - |
173 | | - maybe_slice = lib.maybe_indices_to_slice(indices, len(self)) |
174 | | - |
175 | | - result = NDArrayBackedExtensionIndex.take( |
176 | | - self, indices, axis, allow_fill, fill_value, **kwargs |
177 | | - ) |
178 | | - if isinstance(maybe_slice, slice): |
179 | | - freq = self._data._get_getitem_freq(maybe_slice) |
180 | | - result._data._freq = freq |
181 | | - return result |
182 | | - |
183 | 150 | _can_hold_na = True |
184 | 151 |
|
185 | 152 | _na_value: NaTType = NaT |
@@ -409,102 +376,7 @@ def shift(self: _T, periods: int = 1, freq=None) -> _T: |
409 | 376 | arr = self._data.view() |
410 | 377 | arr._freq = self.freq |
411 | 378 | result = arr._time_shift(periods, freq=freq) |
412 | | - return type(self)(result, name=self.name) |
413 | | - |
414 | | - # -------------------------------------------------------------------- |
415 | | - # List-like Methods |
416 | | - |
417 | | - def _get_delete_freq(self, loc: int | slice | Sequence[int]): |
418 | | - """ |
419 | | - Find the `freq` for self.delete(loc). |
420 | | - """ |
421 | | - freq = None |
422 | | - if is_period_dtype(self.dtype): |
423 | | - freq = self.freq |
424 | | - elif self.freq is not None: |
425 | | - if is_integer(loc): |
426 | | - if loc in (0, -len(self), -1, len(self) - 1): |
427 | | - freq = self.freq |
428 | | - else: |
429 | | - if is_list_like(loc): |
430 | | - # error: Incompatible types in assignment (expression has |
431 | | - # type "Union[slice, ndarray]", variable has type |
432 | | - # "Union[int, slice, Sequence[int]]") |
433 | | - loc = lib.maybe_indices_to_slice( # type: ignore[assignment] |
434 | | - np.asarray(loc, dtype=np.intp), len(self) |
435 | | - ) |
436 | | - if isinstance(loc, slice) and loc.step in (1, None): |
437 | | - if loc.start in (0, None) or loc.stop in (len(self), None): |
438 | | - freq = self.freq |
439 | | - return freq |
440 | | - |
441 | | - def _get_insert_freq(self, loc: int, item): |
442 | | - """ |
443 | | - Find the `freq` for self.insert(loc, item). |
444 | | - """ |
445 | | - value = self._data._validate_scalar(item) |
446 | | - item = self._data._box_func(value) |
447 | | - |
448 | | - freq = None |
449 | | - if is_period_dtype(self.dtype): |
450 | | - freq = self.freq |
451 | | - elif self.freq is not None: |
452 | | - # freq can be preserved on edge cases |
453 | | - if self.size: |
454 | | - if item is NaT: |
455 | | - pass |
456 | | - elif (loc == 0 or loc == -len(self)) and item + self.freq == self[0]: |
457 | | - freq = self.freq |
458 | | - elif (loc == len(self)) and item - self.freq == self[-1]: |
459 | | - freq = self.freq |
460 | | - else: |
461 | | - # Adding a single item to an empty index may preserve freq |
462 | | - if self.freq.is_on_offset(item): |
463 | | - freq = self.freq |
464 | | - return freq |
465 | | - |
466 | | - @doc(NDArrayBackedExtensionIndex.delete) |
467 | | - def delete(self: _T, loc) -> _T: |
468 | | - result = super().delete(loc) |
469 | | - result._data._freq = self._get_delete_freq(loc) |
470 | | - return result |
471 | | - |
472 | | - @doc(NDArrayBackedExtensionIndex.insert) |
473 | | - def insert(self, loc: int, item): |
474 | | - result = super().insert(loc, item) |
475 | | - if isinstance(result, type(self)): |
476 | | - # i.e. parent class method did not cast |
477 | | - result._data._freq = self._get_insert_freq(loc, item) |
478 | | - return result |
479 | | - |
480 | | - # -------------------------------------------------------------------- |
481 | | - # Join/Set Methods |
482 | | - |
483 | | - def _get_join_freq(self, other): |
484 | | - """ |
485 | | - Get the freq to attach to the result of a join operation. |
486 | | - """ |
487 | | - if is_period_dtype(self.dtype): |
488 | | - freq = self.freq |
489 | | - else: |
490 | | - self = cast(DatetimeTimedeltaMixin, self) |
491 | | - freq = self.freq if self._can_fast_union(other) else None |
492 | | - return freq |
493 | | - |
494 | | - def _wrap_joined_index(self, joined, other): |
495 | | - assert other.dtype == self.dtype, (other.dtype, self.dtype) |
496 | | - result = super()._wrap_joined_index(joined, other) |
497 | | - result._data._freq = self._get_join_freq(other) |
498 | | - return result |
499 | | - |
500 | | - def _get_engine_target(self) -> np.ndarray: |
501 | | - # engine methods and libjoin methods need dt64/td64 values cast to i8 |
502 | | - return self._data._ndarray.view("i8") |
503 | | - |
504 | | - def _from_join_target(self, result: np.ndarray): |
505 | | - # view e.g. i8 back to M8[ns] |
506 | | - result = result.view(self._data._ndarray.dtype) |
507 | | - return self._data._from_backing_data(result) |
| 379 | + return type(self)._simple_new(result, name=self.name) |
508 | 380 |
|
509 | 381 | # -------------------------------------------------------------------- |
510 | 382 |
|
@@ -552,6 +424,11 @@ def is_type_compatible(self, kind: str) -> bool: |
552 | 424 | ) |
553 | 425 | return kind in self._data._infer_matches |
554 | 426 |
|
| 427 | + @property |
| 428 | + def values(self) -> np.ndarray: |
| 429 | + # NB: For Datetime64TZ this is lossy |
| 430 | + return self._data._ndarray |
| 431 | + |
555 | 432 | # -------------------------------------------------------------------- |
556 | 433 | # Set Operation Methods |
557 | 434 |
|
@@ -702,3 +579,119 @@ def _union(self, other, sort): |
702 | 579 | return result |
703 | 580 | else: |
704 | 581 | return super()._union(other, sort)._with_freq("infer") |
| 582 | + |
| 583 | + # -------------------------------------------------------------------- |
| 584 | + # Join Methods |
| 585 | + |
| 586 | + def _get_join_freq(self, other): |
| 587 | + """ |
| 588 | + Get the freq to attach to the result of a join operation. |
| 589 | + """ |
| 590 | + freq = None |
| 591 | + if self._can_fast_union(other): |
| 592 | + freq = self.freq |
| 593 | + return freq |
| 594 | + |
| 595 | + def _wrap_joined_index(self, joined, other): |
| 596 | + assert other.dtype == self.dtype, (other.dtype, self.dtype) |
| 597 | + result = super()._wrap_joined_index(joined, other) |
| 598 | + result._data._freq = self._get_join_freq(other) |
| 599 | + return result |
| 600 | + |
| 601 | + def _get_engine_target(self) -> np.ndarray: |
| 602 | + # engine methods and libjoin methods need dt64/td64 values cast to i8 |
| 603 | + return self._data._ndarray.view("i8") |
| 604 | + |
| 605 | + def _from_join_target(self, result: np.ndarray): |
| 606 | + # view e.g. i8 back to M8[ns] |
| 607 | + result = result.view(self._data._ndarray.dtype) |
| 608 | + return self._data._from_backing_data(result) |
| 609 | + |
| 610 | + # -------------------------------------------------------------------- |
| 611 | + # List-like Methods |
| 612 | + |
| 613 | + def _get_delete_freq(self, loc: int | slice | Sequence[int]): |
| 614 | + """ |
| 615 | + Find the `freq` for self.delete(loc). |
| 616 | + """ |
| 617 | + freq = None |
| 618 | + if self.freq is not None: |
| 619 | + if is_integer(loc): |
| 620 | + if loc in (0, -len(self), -1, len(self) - 1): |
| 621 | + freq = self.freq |
| 622 | + else: |
| 623 | + if is_list_like(loc): |
| 624 | + # error: Incompatible types in assignment (expression has |
| 625 | + # type "Union[slice, ndarray]", variable has type |
| 626 | + # "Union[int, slice, Sequence[int]]") |
| 627 | + loc = lib.maybe_indices_to_slice( # type: ignore[assignment] |
| 628 | + np.asarray(loc, dtype=np.intp), len(self) |
| 629 | + ) |
| 630 | + if isinstance(loc, slice) and loc.step in (1, None): |
| 631 | + if loc.start in (0, None) or loc.stop in (len(self), None): |
| 632 | + freq = self.freq |
| 633 | + return freq |
| 634 | + |
| 635 | + def _get_insert_freq(self, loc: int, item): |
| 636 | + """ |
| 637 | + Find the `freq` for self.insert(loc, item). |
| 638 | + """ |
| 639 | + value = self._data._validate_scalar(item) |
| 640 | + item = self._data._box_func(value) |
| 641 | + |
| 642 | + freq = None |
| 643 | + if self.freq is not None: |
| 644 | + # freq can be preserved on edge cases |
| 645 | + if self.size: |
| 646 | + if item is NaT: |
| 647 | + pass |
| 648 | + elif (loc == 0 or loc == -len(self)) and item + self.freq == self[0]: |
| 649 | + freq = self.freq |
| 650 | + elif (loc == len(self)) and item - self.freq == self[-1]: |
| 651 | + freq = self.freq |
| 652 | + else: |
| 653 | + # Adding a single item to an empty index may preserve freq |
| 654 | + if self.freq.is_on_offset(item): |
| 655 | + freq = self.freq |
| 656 | + return freq |
| 657 | + |
| 658 | + @doc(NDArrayBackedExtensionIndex.delete) |
| 659 | + def delete(self, loc): |
| 660 | + result = super().delete(loc) |
| 661 | + result._data._freq = self._get_delete_freq(loc) |
| 662 | + return result |
| 663 | + |
| 664 | + @doc(NDArrayBackedExtensionIndex.insert) |
| 665 | + def insert(self, loc: int, item): |
| 666 | + result = super().insert(loc, item) |
| 667 | + if isinstance(result, type(self)): |
| 668 | + # i.e. parent class method did not cast |
| 669 | + result._data._freq = self._get_insert_freq(loc, item) |
| 670 | + return result |
| 671 | + |
| 672 | + # -------------------------------------------------------------------- |
| 673 | + # NDArray-Like Methods |
| 674 | + |
| 675 | + def __array_wrap__(self, result, context=None): |
| 676 | + """ |
| 677 | + Gets called after a ufunc and other functions. |
| 678 | + """ |
| 679 | + out = super().__array_wrap__(result, context=context) |
| 680 | + if isinstance(out, DatetimeTimedeltaMixin) and self.freq is not None: |
| 681 | + out = out._with_freq("infer") |
| 682 | + return out |
| 683 | + |
| 684 | + @Appender(_index_shared_docs["take"] % _index_doc_kwargs) |
| 685 | + def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs): |
| 686 | + nv.validate_take((), kwargs) |
| 687 | + indices = np.asarray(indices, dtype=np.intp) |
| 688 | + |
| 689 | + result = NDArrayBackedExtensionIndex.take( |
| 690 | + self, indices, axis, allow_fill, fill_value, **kwargs |
| 691 | + ) |
| 692 | + |
| 693 | + maybe_slice = lib.maybe_indices_to_slice(indices, len(self)) |
| 694 | + if isinstance(maybe_slice, slice): |
| 695 | + freq = self._data._get_getitem_freq(maybe_slice) |
| 696 | + result._data._freq = freq |
| 697 | + return result |
0 commit comments