TYP: datetimelike arrays (#52334)

* TYP: use Self in Categorical

* TYP: PeriodArray, Categorical

* TYP: datetimelike arrays

Author: jbrockmendel

pandas/core/arrays/categorical.py

@@ -109,7 +109,6 @@
         Shape,
         SortKind,
         npt,
-        type_t,
     )
 
     from pandas import (
@@ -473,8 +472,8 @@ def _internal_fill_value(self) -> int:
     @classmethod
     def _from_sequence(
         cls, scalars, *, dtype: Dtype | None = None, copy: bool = False
-    ) -> Categorical:
-        return Categorical(scalars, dtype=dtype, copy=copy)
+    ) -> Self:
+        return cls(scalars, dtype=dtype, copy=copy)
 
     @overload
     def astype(self, dtype: npt.DTypeLike, copy: bool = ...) -> np.ndarray:
@@ -563,7 +562,7 @@ def to_list(self):
     @classmethod
     def _from_inferred_categories(
         cls, inferred_categories, inferred_codes, dtype, true_values=None
-    ):
+    ) -> Self:
         """
         Construct a Categorical from inferred values.
 
@@ -632,7 +631,7 @@ def _from_inferred_categories(
     @classmethod
     def from_codes(
         cls, codes, categories=None, ordered=None, dtype: Dtype | None = None
-    ) -> Categorical:
+    ) -> Self:
         """
         Make a Categorical type from codes and categories or dtype.
 
@@ -792,7 +791,7 @@ def _set_categories(self, categories, fastpath: bool = False) -> None:
 
         super().__init__(self._ndarray, new_dtype)
 
-    def _set_dtype(self, dtype: CategoricalDtype) -> Categorical:
+    def _set_dtype(self, dtype: CategoricalDtype) -> Self:
         """
         Internal method for directly updating the CategoricalDtype
 
@@ -808,7 +807,7 @@ def _set_dtype(self, dtype: CategoricalDtype) -> Categorical:
         codes = recode_for_categories(self.codes, self.categories, dtype.categories)
         return type(self)(codes, dtype=dtype, fastpath=True)
 
-    def set_ordered(self, value: bool) -> Categorical:
+    def set_ordered(self, value: bool) -> Self:
         """
         Set the ordered attribute to the boolean value.
 
@@ -822,7 +821,7 @@ def set_ordered(self, value: bool) -> Categorical:
         NDArrayBacked.__init__(cat, cat._ndarray, new_dtype)
         return cat
 
-    def as_ordered(self) -> Categorical:
+    def as_ordered(self) -> Self:
         """
         Set the Categorical to be ordered.
 
@@ -833,7 +832,7 @@ def as_ordered(self) -> Categorical:
         """
         return self.set_ordered(True)
 
-    def as_unordered(self) -> Categorical:
+    def as_unordered(self) -> Self:
         """
         Set the Categorical to be unordered.
 
@@ -912,7 +911,7 @@ def set_categories(self, new_categories, ordered=None, rename: bool = False):
         NDArrayBacked.__init__(cat, codes, new_dtype)
         return cat
 
-    def rename_categories(self, new_categories) -> Categorical:
+    def rename_categories(self, new_categories) -> Self:
         """
         Rename categories.
 
@@ -984,7 +983,7 @@ def rename_categories(self, new_categories) -> Categorical:
         cat._set_categories(new_categories)
         return cat
 
-    def reorder_categories(self, new_categories, ordered=None):
+    def reorder_categories(self, new_categories, ordered=None) -> Self:
         """
         Reorder categories as specified in new_categories.
 
@@ -1027,7 +1026,7 @@ def reorder_categories(self, new_categories, ordered=None):
             )
         return self.set_categories(new_categories, ordered=ordered)
 
-    def add_categories(self, new_categories) -> Categorical:
+    def add_categories(self, new_categories) -> Self:
         """
         Add new categories.
 
@@ -1096,7 +1095,7 @@ def add_categories(self, new_categories) -> Categorical:
         NDArrayBacked.__init__(cat, codes, new_dtype)
         return cat
 
-    def remove_categories(self, removals):
+    def remove_categories(self, removals) -> Self:
         """
         Remove the specified categories.
 
@@ -1152,7 +1151,7 @@ def remove_categories(self, removals):
 
         return self.set_categories(new_categories, ordered=self.ordered, rename=False)
 
-    def remove_unused_categories(self) -> Categorical:
+    def remove_unused_categories(self) -> Self:
         """
         Remove categories which are not used.
 
@@ -1454,7 +1453,7 @@ def memory_usage(self, deep: bool = False) -> int:
         """
         return self._codes.nbytes + self.dtype.categories.memory_usage(deep=deep)
 
-    def isna(self) -> np.ndarray:
+    def isna(self) -> npt.NDArray[np.bool_]:
         """
         Detect missing values
 
@@ -1475,7 +1474,7 @@ def isna(self) -> np.ndarray:
 
     isnull = isna
 
-    def notna(self) -> np.ndarray:
+    def notna(self) -> npt.NDArray[np.bool_]:
         """
         Inverse of isna
 
@@ -1544,8 +1543,8 @@ def value_counts(self, dropna: bool = True) -> Series:
     # "ExtensionDtype"
     @classmethod
     def _empty(  # type: ignore[override]
-        cls: type_t[Categorical], shape: Shape, dtype: CategoricalDtype
-    ) -> Categorical:
+        cls, shape: Shape, dtype: CategoricalDtype
+    ) -> Self:
         """
         Analogous to np.empty(shape, dtype=dtype)
 
@@ -1651,7 +1650,7 @@ def sort_values(
         inplace: Literal[False] = ...,
         ascending: bool = ...,
         na_position: str = ...,
-    ) -> Categorical:
+    ) -> Self:
         ...
 
     @overload
@@ -1666,7 +1665,7 @@ def sort_values(
         inplace: bool = False,
         ascending: bool = True,
         na_position: str = "last",
-    ) -> Categorical | None:
+    ) -> Self | None:
         """
         Sort the Categorical by category value returning a new
         Categorical by default.
@@ -1769,7 +1768,7 @@ def _rank(
             pct=pct,
         )
 
-    def _values_for_rank(self):
+    def _values_for_rank(self) -> np.ndarray:
         """
         For correctly ranking ordered categorical data. See GH#15420
 

pandas/core/arrays/datetimelike.py

@@ -1094,6 +1094,7 @@ def _sub_datetimelike(self, other: Timestamp | DatetimeArray) -> TimedeltaArray:
         res_m8 = res_values.view(f"timedelta64[{self.unit}]")
 
         new_freq = self._get_arithmetic_result_freq(other)
+        new_freq = cast("Tick | None", new_freq)
         return TimedeltaArray._simple_new(res_m8, dtype=res_m8.dtype, freq=new_freq)
 
     @final
@@ -1161,7 +1162,12 @@ def _add_timedeltalike(self, other: Timedelta | TimedeltaArray):
 
         new_freq = self._get_arithmetic_result_freq(other)
 
-        return type(self)._simple_new(res_values, dtype=self.dtype, freq=new_freq)
+        # error: Argument "dtype" to "_simple_new" of "DatetimeArray" has
+        # incompatible type "Union[dtype[datetime64], DatetimeTZDtype,
+        # dtype[timedelta64]]"; expected "Union[dtype[datetime64], DatetimeTZDtype]"
+        return type(self)._simple_new(
+            res_values, dtype=self.dtype, freq=new_freq  # type: ignore[arg-type]
+        )
 
     @final
     def _add_nat(self):
@@ -1179,7 +1185,12 @@ def _add_nat(self):
         result = np.empty(self.shape, dtype=np.int64)
         result.fill(iNaT)
         result = result.view(self._ndarray.dtype)  # preserve reso
-        return type(self)._simple_new(result, dtype=self.dtype, freq=None)
+        # error: Argument "dtype" to "_simple_new" of "DatetimeArray" has
+        # incompatible type "Union[dtype[timedelta64], dtype[datetime64],
+        # DatetimeTZDtype]"; expected "Union[dtype[datetime64], DatetimeTZDtype]"
+        return type(self)._simple_new(
+            result, dtype=self.dtype, freq=None  # type: ignore[arg-type]
+        )
 
     @final
     def _sub_nat(self):

pandas/core/arrays/datetimes.py

@@ -81,6 +81,7 @@
     from pandas._typing import (
         DateTimeErrorChoices,
         IntervalClosedType,
+        Self,
         TimeAmbiguous,
         TimeNonexistent,
         npt,
@@ -92,7 +93,9 @@
 _midnight = time(0, 0)
 
 
-def tz_to_dtype(tz: tzinfo | None, unit: str = "ns"):
+def tz_to_dtype(
+    tz: tzinfo | None, unit: str = "ns"
+) -> np.dtype[np.datetime64] | DatetimeTZDtype:
     """
     Return a datetime64[ns] dtype appropriate for the given timezone.
 
@@ -254,7 +257,7 @@ def _scalar_type(self) -> type[Timestamp]:
     # -----------------------------------------------------------------
     # Constructors
 
-    _dtype: np.dtype | DatetimeTZDtype
+    _dtype: np.dtype[np.datetime64] | DatetimeTZDtype
     _freq: BaseOffset | None = None
     _default_dtype = DT64NS_DTYPE  # used in TimeLikeOps.__init__
 
@@ -269,10 +272,10 @@ def _validate_dtype(cls, values, dtype):
     @classmethod
     def _simple_new(  # type: ignore[override]
         cls,
-        values: np.ndarray,
+        values: npt.NDArray[np.datetime64],
         freq: BaseOffset | None = None,
-        dtype=DT64NS_DTYPE,
-    ) -> DatetimeArray:
+        dtype: np.dtype[np.datetime64] | DatetimeTZDtype = DT64NS_DTYPE,
+    ) -> Self:
         assert isinstance(values, np.ndarray)
         assert dtype.kind == "M"
         if isinstance(dtype, np.dtype):
@@ -386,7 +389,7 @@ def _generate_range(  # type: ignore[override]
         inclusive: IntervalClosedType = "both",
         *,
         unit: str | None = None,
-    ) -> DatetimeArray:
+    ) -> Self:
         periods = dtl.validate_periods(periods)
         if freq is None and any(x is None for x in [periods, start, end]):
             raise ValueError("Must provide freq argument if no data is supplied")
@@ -535,7 +538,7 @@ def _box_func(self, x: np.datetime64) -> Timestamp | NaTType:
     # error: Return type "Union[dtype, DatetimeTZDtype]" of "dtype"
     # incompatible with return type "ExtensionDtype" in supertype
     # "ExtensionArray"
-    def dtype(self) -> np.dtype | DatetimeTZDtype:  # type: ignore[override]
+    def dtype(self) -> np.dtype[np.datetime64] | DatetimeTZDtype:  # type: ignore[override]  # noqa:E501
         """
         The dtype for the DatetimeArray.
 
@@ -753,7 +756,7 @@ def _assert_tzawareness_compat(self, other) -> None:
     # -----------------------------------------------------------------
     # Arithmetic Methods
 
-    def _add_offset(self, offset) -> DatetimeArray:
+    def _add_offset(self, offset) -> Self:
         assert not isinstance(offset, Tick)
 
         if self.tz is not None:
@@ -776,7 +779,7 @@ def _add_offset(self, offset) -> DatetimeArray:
                 return result.tz_localize(self.tz)
 
         else:
-            result = DatetimeArray._simple_new(result, dtype=result.dtype)
+            result = type(self)._simple_new(result, dtype=result.dtype)
             if self.tz is not None:
                 result = result.tz_localize(self.tz)
 
@@ -797,7 +800,7 @@ def _local_timestamps(self) -> npt.NDArray[np.int64]:
             return self.asi8
         return tz_convert_from_utc(self.asi8, self.tz, reso=self._creso)
 
-    def tz_convert(self, tz) -> DatetimeArray:
+    def tz_convert(self, tz) -> Self:
         """
         Convert tz-aware Datetime Array/Index from one time zone to another.
 
@@ -879,7 +882,7 @@ def tz_localize(
         tz,
         ambiguous: TimeAmbiguous = "raise",
         nonexistent: TimeNonexistent = "raise",
-    ) -> DatetimeArray:
+    ) -> Self:
         """
         Localize tz-naive Datetime Array/Index to tz-aware Datetime Array/Index.
 
@@ -1045,18 +1048,18 @@ def tz_localize(
                 nonexistent=nonexistent,
                 creso=self._creso,
             )
-        new_dates = new_dates.view(f"M8[{self.unit}]")
+        new_dates_dt64 = new_dates.view(f"M8[{self.unit}]")
         dtype = tz_to_dtype(tz, unit=self.unit)
 
         freq = None
-        if timezones.is_utc(tz) or (len(self) == 1 and not isna(new_dates[0])):
+        if timezones.is_utc(tz) or (len(self) == 1 and not isna(new_dates_dt64[0])):
             # we can preserve freq
             # TODO: Also for fixed-offsets
             freq = self.freq
         elif tz is None and self.tz is None:
             # no-op
             freq = self.freq
-        return self._simple_new(new_dates, dtype=dtype, freq=freq)
+        return self._simple_new(new_dates_dt64, dtype=dtype, freq=freq)
 
     # ----------------------------------------------------------------
     # Conversion Methods - Vectorized analogues of Timestamp methods
@@ -1071,7 +1074,7 @@ def to_pydatetime(self) -> npt.NDArray[np.object_]:
         """
         return ints_to_pydatetime(self.asi8, tz=self.tz, reso=self._creso)
 
-    def normalize(self) -> DatetimeArray:
+    def normalize(self) -> Self:
         """
         Convert times to midnight.