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

Author: topper-123

pandas/core/internals/array_manager.py

@@ -8,7 +8,6 @@
     Callable,
     Hashable,
     Literal,
-    TypeVar,
 )
 
 import numpy as np
@@ -23,6 +22,7 @@
     AxisInt,
     DtypeObj,
     QuantileInterpolation,
+    Self,
     npt,
 )
 from pandas.util._validators import validate_bool_kwarg
@@ -93,8 +93,6 @@
     to_native_types,
 )
 
-T = TypeVar("T", bound="BaseArrayManager")
-
 
 class BaseArrayManager(DataManager):
     """
@@ -129,7 +127,7 @@ def __init__(
     ) -> None:
         raise NotImplementedError
 
-    def make_empty(self: T, axes=None) -> T:
+    def make_empty(self, axes=None) -> Self:
         """Return an empty ArrayManager with the items axis of len 0 (no columns)"""
         if axes is None:
             axes = [self.axes[1:], Index([])]
@@ -193,11 +191,11 @@ def __repr__(self) -> str:
         return output
 
     def apply(
-        self: T,
+        self,
         f,
         align_keys: list[str] | None = None,
         **kwargs,
-    ) -> T:
+    ) -> Self:
         """
         Iterate over the arrays, collect and create a new ArrayManager.
 
@@ -255,8 +253,8 @@ def apply(
         return type(self)(result_arrays, new_axes)  # type: ignore[arg-type]
 
     def apply_with_block(
-        self: T, f, align_keys=None, swap_axis: bool = True, **kwargs
-    ) -> T:
+        self, f, align_keys=None, swap_axis: bool = True, **kwargs
+    ) -> Self:
         # switch axis to follow BlockManager logic
         if swap_axis and "axis" in kwargs and self.ndim == 2:
             kwargs["axis"] = 1 if kwargs["axis"] == 0 else 0
@@ -309,7 +307,7 @@ def apply_with_block(
 
         return type(self)(result_arrays, self._axes)
 
-    def where(self: T, other, cond, align: bool) -> T:
+    def where(self, other, cond, align: bool) -> Self:
         if align:
             align_keys = ["other", "cond"]
         else:
@@ -323,10 +321,10 @@ def where(self: T, other, cond, align: bool) -> T:
             cond=cond,
         )
 
-    def setitem(self: T, indexer, value) -> T:
+    def setitem(self, indexer, value) -> Self:
         return self.apply_with_block("setitem", indexer=indexer, value=value)
 
-    def putmask(self: T, mask, new, align: bool = True) -> T:
+    def putmask(self, mask, new, align: bool = True) -> Self:
         if align:
             align_keys = ["new", "mask"]
         else:
@@ -340,14 +338,14 @@ def putmask(self: T, mask, new, align: bool = True) -> T:
             new=new,
         )
 
-    def diff(self: T, n: int, axis: AxisInt) -> T:
+    def diff(self, n: int, axis: AxisInt) -> Self:
         assert self.ndim == 2 and axis == 0  # caller ensures
         return self.apply(algos.diff, n=n, axis=axis)
 
-    def interpolate(self: T, **kwargs) -> T:
+    def interpolate(self, **kwargs) -> Self:
         return self.apply_with_block("interpolate", swap_axis=False, **kwargs)
 
-    def shift(self: T, periods: int, axis: AxisInt, fill_value) -> T:
+    def shift(self, periods: int, axis: AxisInt, fill_value) -> Self:
         if fill_value is lib.no_default:
             fill_value = None
 
@@ -359,7 +357,7 @@ def shift(self: T, periods: int, axis: AxisInt, fill_value) -> T:
             "shift", periods=periods, axis=axis, fill_value=fill_value
         )
 
-    def fillna(self: T, value, limit, inplace: bool, downcast) -> T:
+    def fillna(self, value, limit, inplace: bool, downcast) -> Self:
         if limit is not None:
             # Do this validation even if we go through one of the no-op paths
             limit = libalgos.validate_limit(None, limit=limit)
@@ -368,13 +366,13 @@ def fillna(self: T, value, limit, inplace: bool, downcast) -> T:
             "fillna", value=value, limit=limit, inplace=inplace, downcast=downcast
         )
 
-    def astype(self: T, dtype, copy: bool | None = False, errors: str = "raise") -> T:
+    def astype(self, dtype, copy: bool | None = False, errors: str = "raise") -> Self:
         if copy is None:
             copy = True
 
         return self.apply(astype_array_safe, dtype=dtype, copy=copy, errors=errors)
 
-    def convert(self: T, copy: bool | None) -> T:
+    def convert(self, copy: bool | None) -> Self:
         if copy is None:
             copy = True
 
@@ -397,10 +395,10 @@ def _convert(arr):
 
         return self.apply(_convert)
 
-    def replace_regex(self: T, **kwargs) -> T:
+    def replace_regex(self, **kwargs) -> Self:
         return self.apply_with_block("_replace_regex", **kwargs)
 
-    def replace(self: T, to_replace, value, inplace: bool) -> T:
+    def replace(self, to_replace, value, inplace: bool) -> Self:
         inplace = validate_bool_kwarg(inplace, "inplace")
         assert np.ndim(value) == 0, value
         # TODO "replace" is right now implemented on the blocks, we should move
@@ -410,12 +408,12 @@ def replace(self: T, to_replace, value, inplace: bool) -> T:
         )
 
     def replace_list(
-        self: T,
+        self,
         src_list: list[Any],
         dest_list: list[Any],
         inplace: bool = False,
         regex: bool = False,
-    ) -> T:
+    ) -> Self:
         """do a list replace"""
         inplace = validate_bool_kwarg(inplace, "inplace")
 
@@ -427,7 +425,7 @@ def replace_list(
             regex=regex,
         )
 
-    def to_native_types(self: T, **kwargs) -> T:
+    def to_native_types(self, **kwargs) -> Self:
         return self.apply(to_native_types, **kwargs)
 
     @property
@@ -453,7 +451,7 @@ def is_view(self) -> bool:
     def is_single_block(self) -> bool:
         return len(self.arrays) == 1
 
-    def _get_data_subset(self: T, predicate: Callable) -> T:
+    def _get_data_subset(self, predicate: Callable) -> Self:
         indices = [i for i, arr in enumerate(self.arrays) if predicate(arr)]
         arrays = [self.arrays[i] for i in indices]
         # TODO copy?
@@ -464,7 +462,7 @@ def _get_data_subset(self: T, predicate: Callable) -> T:
         new_axes = [self._axes[0], new_cols]
         return type(self)(arrays, new_axes, verify_integrity=False)
 
-    def get_bool_data(self: T, copy: bool = False) -> T:
+    def get_bool_data(self, copy: bool = False) -> Self:
         """
         Select columns that are bool-dtype and object-dtype columns that are all-bool.
 
@@ -475,7 +473,7 @@ def get_bool_data(self: T, copy: bool = False) -> T:
         """
         return self._get_data_subset(lambda x: x.dtype == np.dtype(bool))
 
-    def get_numeric_data(self: T, copy: bool = False) -> T:
+    def get_numeric_data(self, copy: bool = False) -> Self:
         """
         Select columns that have a numeric dtype.
 
@@ -489,7 +487,7 @@ def get_numeric_data(self: T, copy: bool = False) -> T:
             or getattr(arr.dtype, "_is_numeric", False)
         )
 
-    def copy(self: T, deep: bool | Literal["all"] | None = True) -> T:
+    def copy(self, deep: bool | Literal["all"] | None = True) -> Self:
         """
         Make deep or shallow copy of ArrayManager
 
@@ -526,7 +524,7 @@ def copy_func(ax):
         return type(self)(new_arrays, new_axes, verify_integrity=False)
 
     def reindex_indexer(
-        self: T,
+        self,
         new_axis,
         indexer,
         axis: AxisInt,
@@ -537,7 +535,7 @@ def reindex_indexer(
         only_slice: bool = False,
         # ArrayManager specific keywords
         use_na_proxy: bool = False,
-    ) -> T:
+    ) -> Self:
         axis = self._normalize_axis(axis)
         return self._reindex_indexer(
             new_axis,
@@ -550,15 +548,15 @@ def reindex_indexer(
         )
 
     def _reindex_indexer(
-        self: T,
+        self,
         new_axis,
         indexer: npt.NDArray[np.intp] | None,
         axis: AxisInt,
         fill_value=None,
         allow_dups: bool = False,
         copy: bool | None = True,
         use_na_proxy: bool = False,
-    ) -> T:
+    ) -> Self:
         """
         Parameters
         ----------
@@ -629,12 +627,12 @@ def _reindex_indexer(
         return type(self)(new_arrays, new_axes, verify_integrity=False)
 
     def take(
-        self: T,
+        self,
         indexer,
         axis: AxisInt = 1,
         verify: bool = True,
         convert_indices: bool = True,
-    ) -> T:
+    ) -> Self:
         """
         Take items along any axis.
         """
@@ -926,7 +924,7 @@ def idelete(self, indexer) -> ArrayManager:
     # --------------------------------------------------------------------
     # Array-wise Operation
 
-    def grouped_reduce(self: T, func: Callable) -> T:
+    def grouped_reduce(self, func: Callable) -> Self:
         """
         Apply grouped reduction function columnwise, returning a new ArrayManager.
 
@@ -965,7 +963,7 @@ def grouped_reduce(self: T, func: Callable) -> T:
         # expected "List[Union[ndarray, ExtensionArray]]"
         return type(self)(result_arrays, [index, columns])  # type: ignore[arg-type]
 
-    def reduce(self: T, func: Callable) -> T:
+    def reduce(self, func: Callable) -> Self:
         """
         Apply reduction function column-wise, returning a single-row ArrayManager.
 

pandas/core/internals/base.py

@@ -6,7 +6,6 @@
 
 from typing import (
     Literal,
-    TypeVar,
     final,
 )
 
@@ -16,6 +15,7 @@
     ArrayLike,
     AxisInt,
     DtypeObj,
+    Self,
     Shape,
 )
 from pandas.errors import AbstractMethodError
@@ -31,8 +31,6 @@
     default_index,
 )
 
-T = TypeVar("T", bound="DataManager")
-
 
 class DataManager(PandasObject):
     # TODO share more methods/attributes
@@ -73,25 +71,25 @@ def _validate_set_axis(self, axis: AxisInt, new_labels: Index) -> None:
             )
 
     def reindex_indexer(
-        self: T,
+        self,
         new_axis,
         indexer,
         axis: AxisInt,
         fill_value=None,
         allow_dups: bool = False,
         copy: bool = True,
         only_slice: bool = False,
-    ) -> T:
+    ) -> Self:
         raise AbstractMethodError(self)
 
     @final
     def reindex_axis(
-        self: T,
+        self,
         new_index: Index,
         axis: AxisInt,
         fill_value=None,
         only_slice: bool = False,
-    ) -> T:
+    ) -> Self:
         """
         Conform data manager to new index.
         """
@@ -106,7 +104,7 @@ def reindex_axis(
             only_slice=only_slice,
         )
 
-    def _equal_values(self: T, other: T) -> bool:
+    def _equal_values(self, other: Self) -> bool:
         """
         To be implemented by the subclasses. Only check the column values
         assuming shape and indexes have already been checked.
@@ -130,15 +128,15 @@ def equals(self, other: object) -> bool:
         return self._equal_values(other)
 
     def apply(
-        self: T,
+        self,
         f,
         align_keys: list[str] | None = None,
         **kwargs,
-    ) -> T:
+    ) -> Self:
         raise AbstractMethodError(self)
 
     @final
-    def isna(self: T, func) -> T:
+    def isna(self, func) -> Self:
         return self.apply("apply", func=func)
 
     # --------------------------------------------------------------------
@@ -147,7 +145,7 @@ def isna(self: T, func) -> T:
     def is_consolidated(self) -> bool:
         return True
 
-    def consolidate(self: T) -> T:
+    def consolidate(self) -> Self:
         return self
 
     def _consolidate_inplace(self) -> None: