TYP: Use Self for type checking (pandas/core/arrays/)

Author: topper-123

pandas/_typing.py

@@ -6,6 +6,7 @@
     tzinfo,
 )
 from os import PathLike
+import sys
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -83,8 +84,13 @@
     # Name "npt._ArrayLikeInt_co" is not defined  [name-defined]
     NumpySorter = Optional[npt._ArrayLikeInt_co]  # type: ignore[name-defined]
 
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self  # pyright: reportUnusedImport = false
 else:
     npt: Any = None
+    Self: Any = None
 
 HashableT = TypeVar("HashableT", bound=Hashable)
 

pandas/core/arrays/_mixins.py

@@ -6,7 +6,6 @@
     Any,
     Literal,
     Sequence,
-    TypeVar,
     cast,
     overload,
 )
@@ -23,11 +22,11 @@
     PositionalIndexer2D,
     PositionalIndexerTuple,
     ScalarIndexer,
+    Self,
     SequenceIndexer,
     Shape,
     TakeIndexer,
     npt,
-    type_t,
 )
 from pandas.errors import AbstractMethodError
 from pandas.util._decorators import doc
@@ -61,10 +60,6 @@
 from pandas.core.indexers import check_array_indexer
 from pandas.core.sorting import nargminmax
 
-NDArrayBackedExtensionArrayT = TypeVar(
-    "NDArrayBackedExtensionArrayT", bound="NDArrayBackedExtensionArray"
-)
-
 if TYPE_CHECKING:
     from pandas._typing import (
         NumpySorter,
@@ -153,13 +148,13 @@ def view(self, dtype: Dtype | None = None) -> ArrayLike:
         return arr.view(dtype=dtype)  # type: ignore[arg-type]
 
     def take(
-        self: NDArrayBackedExtensionArrayT,
+        self,
         indices: TakeIndexer,
         *,
         allow_fill: bool = False,
         fill_value: Any = None,
         axis: AxisInt = 0,
-    ) -> NDArrayBackedExtensionArrayT:
+    ) -> Self:
         if allow_fill:
             fill_value = self._validate_scalar(fill_value)
 
@@ -208,17 +203,17 @@ def argmax(self, axis: AxisInt = 0, skipna: bool = True):  # type: ignore[overri
             raise NotImplementedError
         return nargminmax(self, "argmax", axis=axis)
 
-    def unique(self: NDArrayBackedExtensionArrayT) -> NDArrayBackedExtensionArrayT:
+    def unique(self) -> Self:
         new_data = unique(self._ndarray)
         return self._from_backing_data(new_data)
 
     @classmethod
     @doc(ExtensionArray._concat_same_type)
     def _concat_same_type(
-        cls: type[NDArrayBackedExtensionArrayT],
-        to_concat: Sequence[NDArrayBackedExtensionArrayT],
+        cls,
+        to_concat: Sequence[Self],
         axis: AxisInt = 0,
-    ) -> NDArrayBackedExtensionArrayT:
+    ) -> Self:
         dtypes = {str(x.dtype) for x in to_concat}
         if len(dtypes) != 1:
             raise ValueError("to_concat must have the same dtype (tz)", dtypes)
@@ -258,15 +253,15 @@ def __getitem__(self, key: ScalarIndexer) -> Any:
 
     @overload
     def __getitem__(
-        self: NDArrayBackedExtensionArrayT,
+        self,
         key: SequenceIndexer | PositionalIndexerTuple,
-    ) -> NDArrayBackedExtensionArrayT:
+    ) -> Self:
         ...
 
     def __getitem__(
-        self: NDArrayBackedExtensionArrayT,
+        self,
         key: PositionalIndexer2D,
-    ) -> NDArrayBackedExtensionArrayT | Any:
+    ) -> Self | Any:
         if lib.is_integer(key):
             # fast-path
             result = self._ndarray[key]
@@ -293,9 +288,7 @@ def _fill_mask_inplace(
         func(self._ndarray.T, limit=limit, mask=mask.T)
 
     @doc(ExtensionArray.fillna)
-    def fillna(
-        self: NDArrayBackedExtensionArrayT, value=None, method=None, limit=None
-    ) -> NDArrayBackedExtensionArrayT:
+    def fillna(self, value=None, method=None, limit=None) -> Self:
         value, method = validate_fillna_kwargs(
             value, method, validate_scalar_dict_value=False
         )
@@ -359,9 +352,7 @@ def _putmask(self, mask: npt.NDArray[np.bool_], value) -> None:
 
         np.putmask(self._ndarray, mask, value)
 
-    def _where(
-        self: NDArrayBackedExtensionArrayT, mask: npt.NDArray[np.bool_], value
-    ) -> NDArrayBackedExtensionArrayT:
+    def _where(self: Self, mask: npt.NDArray[np.bool_], value) -> Self:
         """
         Analogue to np.where(mask, self, value)
 
@@ -383,9 +374,7 @@ def _where(
     # ------------------------------------------------------------------------
     # Index compat methods
 
-    def insert(
-        self: NDArrayBackedExtensionArrayT, loc: int, item
-    ) -> NDArrayBackedExtensionArrayT:
+    def insert(self, loc: int, item) -> Self:
         """
         Make new ExtensionArray inserting new item at location. Follows
         Python list.append semantics for negative values.
@@ -451,10 +440,10 @@ def value_counts(self, dropna: bool = True) -> Series:
         return Series(result._values, index=index, name=result.name)
 
     def _quantile(
-        self: NDArrayBackedExtensionArrayT,
+        self,
         qs: npt.NDArray[np.float64],
         interpolation: str,
-    ) -> NDArrayBackedExtensionArrayT:
+    ) -> Self:
         # TODO: disable for Categorical if not ordered?
 
         mask = np.asarray(self.isna())
@@ -478,9 +467,7 @@ def _cast_quantile_result(self, res_values: np.ndarray) -> np.ndarray:
     # numpy-like methods
 
     @classmethod
-    def _empty(
-        cls: type_t[NDArrayBackedExtensionArrayT], shape: Shape, dtype: ExtensionDtype
-    ) -> NDArrayBackedExtensionArrayT:
+    def _empty(cls, shape: Shape, dtype: ExtensionDtype) -> Self:
         """
         Analogous to np.empty(shape, dtype=dtype)
 

pandas/core/arrays/arrow/array.py

@@ -9,7 +9,6 @@
     Callable,
     Literal,
     Sequence,
-    TypeVar,
     cast,
 )
 
@@ -25,6 +24,7 @@
     NpDtype,
     PositionalIndexer,
     Scalar,
+    Self,
     SortKind,
     TakeIndexer,
     TimeAmbiguous,
@@ -138,8 +138,6 @@ def floordiv_compat(
 
     from pandas import Series
 
-ArrowExtensionArrayT = TypeVar("ArrowExtensionArrayT", bound="ArrowExtensionArray")
-
 
 def get_unit_from_pa_dtype(pa_dtype):
     # https://github.com/pandas-dev/pandas/pull/50998#discussion_r1100344804
@@ -405,16 +403,16 @@ def __array__(self, dtype: NpDtype | None = None) -> np.ndarray:
         """Correctly construct numpy arrays when passed to `np.asarray()`."""
         return self.to_numpy(dtype=dtype)
 
-    def __invert__(self: ArrowExtensionArrayT) -> ArrowExtensionArrayT:
+    def __invert__(self) -> Self:
         return type(self)(pc.invert(self._data))
 
-    def __neg__(self: ArrowExtensionArrayT) -> ArrowExtensionArrayT:
+    def __neg__(self) -> Self:
         return type(self)(pc.negate_checked(self._data))
 
-    def __pos__(self: ArrowExtensionArrayT) -> ArrowExtensionArrayT:
+    def __pos__(self) -> Self:
         return type(self)(self._data)
 
-    def __abs__(self: ArrowExtensionArrayT) -> ArrowExtensionArrayT:
+    def __abs__(self) -> Self:
         return type(self)(pc.abs_checked(self._data))
 
     # GH 42600: __getstate__/__setstate__ not necessary once
@@ -606,7 +604,7 @@ def argmin(self, skipna: bool = True) -> int:
     def argmax(self, skipna: bool = True) -> int:
         return self._argmin_max(skipna, "max")
 
-    def copy(self: ArrowExtensionArrayT) -> ArrowExtensionArrayT:
+    def copy(self) -> Self:
         """
         Return a shallow copy of the array.
 
@@ -618,7 +616,7 @@ def copy(self: ArrowExtensionArrayT) -> ArrowExtensionArrayT:
         """
         return type(self)(self._data)
 
-    def dropna(self: ArrowExtensionArrayT) -> ArrowExtensionArrayT:
+    def dropna(self) -> Self:
         """
         Return ArrowExtensionArray without NA values.
 
@@ -630,11 +628,11 @@ def dropna(self: ArrowExtensionArrayT) -> ArrowExtensionArrayT:
 
     @doc(ExtensionArray.fillna)
     def fillna(
-        self: ArrowExtensionArrayT,
+        self,
         value: object | ArrayLike | None = None,
         method: FillnaOptions | None = None,
         limit: int | None = None,
-    ) -> ArrowExtensionArrayT:
+    ) -> Self:
         value, method = validate_fillna_kwargs(value, method)
 
         if limit is not None:
@@ -751,9 +749,7 @@ def reshape(self, *args, **kwargs):
             f"as backed by a 1D pyarrow.ChunkedArray."
         )
 
-    def round(
-        self: ArrowExtensionArrayT, decimals: int = 0, *args, **kwargs
-    ) -> ArrowExtensionArrayT:
+    def round(self, decimals: int = 0, *args, **kwargs) -> Self:
         """
         Round each value in the array a to the given number of decimals.
 
@@ -926,7 +922,7 @@ def to_numpy(
             result[self.isna()] = na_value
         return result
 
-    def unique(self: ArrowExtensionArrayT) -> ArrowExtensionArrayT:
+    def unique(self) -> Self:
         """
         Compute the ArrowExtensionArray of unique values.
 
@@ -998,9 +994,7 @@ def value_counts(self, dropna: bool = True) -> Series:
         return Series(counts, index=index, name="count").astype("Int64")
 
     @classmethod
-    def _concat_same_type(
-        cls: type[ArrowExtensionArrayT], to_concat
-    ) -> ArrowExtensionArrayT:
+    def _concat_same_type(cls, to_concat) -> Self:
         """
         Concatenate multiple ArrowExtensionArrays.
 
@@ -1321,9 +1315,7 @@ def _rank(
 
         return type(self)(result)
 
-    def _quantile(
-        self: ArrowExtensionArrayT, qs: npt.NDArray[np.float64], interpolation: str
-    ) -> ArrowExtensionArrayT:
+    def _quantile(self, qs: npt.NDArray[np.float64], interpolation: str) -> Self:
         """
         Compute the quantiles of self for each quantile in `qs`.
 
@@ -1360,7 +1352,7 @@ def _quantile(
 
         return type(self)(result)
 
-    def _mode(self: ArrowExtensionArrayT, dropna: bool = True) -> ArrowExtensionArrayT:
+    def _mode(self, dropna: bool = True) -> Self:
         """
         Returns the mode(s) of the ExtensionArray.