Expunge most uses of TypeVar(bound="Foo") (#13346)

With #13336, we can use the PEP674 `Self` type.

Authors:
  - Lawrence Mitchell (https://github.com/wence-)

Approvers:
  - Bradley Dice (https://github.com/bdice)

URL: #13346

python/cudf/cudf/_lib/column.pyi

@@ -4,12 +4,12 @@ from __future__ import annotations
 
 from typing import Dict, Optional, Tuple, TypeVar
 
+from typing_extensions import Self
+
 from cudf._typing import Dtype, DtypeObj, ScalarLike
 from cudf.core.buffer import Buffer
 from cudf.core.column import ColumnBase
 
-T = TypeVar("T")
-
 class Column:
     _data: Optional[Buffer]
     _mask: Optional[Buffer]
@@ -56,7 +56,7 @@ class Column:
     @property
     def mask_ptr(self) -> int: ...
     def set_base_mask(self, value: Optional[Buffer]) -> None: ...
-    def set_mask(self: T, value: Optional[Buffer]) -> T: ...
+    def set_mask(self, value: Optional[Buffer]) -> Self: ...
     @property
     def null_count(self) -> int: ...
     @property
@@ -68,7 +68,8 @@ class Column:
     def set_base_children(self, value: Tuple[ColumnBase, ...]) -> None: ...
     def _mimic_inplace(
         self, other_col: ColumnBase, inplace=False
-    ) -> Optional[ColumnBase]: ...
+    ) -> Optional[Self]: ...
+
     # TODO: The val parameter should be Scalar, not ScalarLike
     @staticmethod
     def from_scalar(val: ScalarLike, size: int) -> ColumnBase: ...

python/cudf/cudf/core/_base_index.py

@@ -5,9 +5,10 @@
 import pickle
 import warnings
 from functools import cached_property
-from typing import Any, Set, TypeVar
+from typing import Any, Set
 
 import pandas as pd
+from typing_extensions import Self
 
 import cudf
 from cudf._lib.copying import _gather_map_is_valid, gather
@@ -62,8 +63,6 @@
 Float64Index([1.0, 2.0, 3.0], dtype='float64')
 """
 
-BaseIndexT = TypeVar("BaseIndexT", bound="BaseIndex")
-
 
 class BaseIndex(Serializable):
     """Base class for all cudf Index types."""
@@ -101,8 +100,8 @@ def __contains__(self, item):
         return item in self._values
 
     def _copy_type_metadata(
-        self: BaseIndexT, other: BaseIndexT, *, override_dtypes=None
-    ) -> BaseIndexT:
+        self, other: Self, *, override_dtypes=None
+    ) -> Self:
         raise NotImplementedError
 
     def get_level_values(self, level):
@@ -1451,7 +1450,6 @@ def get_slice_bound(self, label, side, kind=None):
         raise NotImplementedError
 
     def __array_function__(self, func, types, args, kwargs):
-
         # check if the function is implemented for the current type
         cudf_index_module = type(self)
         for submodule in func.__module__.split(".")[1:]:

python/cudf/cudf/core/buffer/buffer.py

@@ -5,18 +5,17 @@
 import math
 import pickle
 from types import SimpleNamespace
-from typing import Any, Dict, Mapping, Optional, Sequence, Tuple, Type, TypeVar
+from typing import Any, Dict, Mapping, Optional, Sequence, Tuple
 
 import numpy
+from typing_extensions import Self
 
 import rmm
 
 import cudf
 from cudf.core.abc import Serializable
 from cudf.utils.string import format_bytes
 
-T = TypeVar("T", bound="Buffer")
-
 
 def host_memory_allocation(nbytes: int) -> memoryview:
     """Allocate host memory using NumPy
@@ -108,7 +107,7 @@ def __init__(self):
         )
 
     @classmethod
-    def _from_device_memory(cls: Type[T], data: Any) -> T:
+    def _from_device_memory(cls, data: Any) -> Self:
         """Create a Buffer from an object exposing `__cuda_array_interface__`.
 
         No data is being copied.
@@ -139,7 +138,7 @@ def _from_device_memory(cls: Type[T], data: Any) -> T:
         return ret
 
     @classmethod
-    def _from_host_memory(cls: Type[T], data: Any) -> T:
+    def _from_host_memory(cls, data: Any) -> Self:
         """Create a Buffer from a buffer or array like object
 
         Data must implement `__array_interface__`, the buffer protocol, and/or
@@ -310,7 +309,7 @@ def serialize(self) -> Tuple[dict, list]:
         return header, frames
 
     @classmethod
-    def deserialize(cls: Type[T], header: dict, frames: list) -> T:
+    def deserialize(cls, header: dict, frames: list) -> Self:
         """Create an Buffer from a serialized representation.
 
         Parameters

python/cudf/cudf/core/buffer/cow_buffer.py

@@ -4,15 +4,15 @@
 
 import weakref
 from collections import defaultdict
-from typing import Any, DefaultDict, Tuple, Type, TypeVar
+from typing import Any, DefaultDict, Tuple
 from weakref import WeakSet
 
+from typing_extensions import Self
+
 import rmm
 
 from cudf.core.buffer.buffer import Buffer
 
-T = TypeVar("T", bound="CopyOnWriteBuffer")
-
 
 def _keys_cleanup(ptr):
     weak_set_values = CopyOnWriteBuffer._instances[ptr]
@@ -55,9 +55,7 @@ def _finalize_init(self):
         weakref.finalize(self, _keys_cleanup, self._ptr)
 
     @classmethod
-    def _from_device_memory(
-        cls: Type[T], data: Any, *, exposed: bool = False
-    ) -> T:
+    def _from_device_memory(cls, data: Any, *, exposed: bool = False) -> Self:
         """Create a Buffer from an object exposing `__cuda_array_interface__`.
 
         No data is being copied.
@@ -82,7 +80,7 @@ def _from_device_memory(
         return ret
 
     @classmethod
-    def _from_host_memory(cls: Type[T], data: Any) -> T:
+    def _from_host_memory(cls, data: Any) -> Self:
         ret = super()._from_host_memory(data)
         ret._finalize_init()
         return ret

python/cudf/cudf/core/buffer/spillable_buffer.py

@@ -7,18 +7,10 @@
 import time
 import weakref
 from threading import RLock
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Dict,
-    List,
-    Optional,
-    Tuple,
-    Type,
-    TypeVar,
-)
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
 import numpy
+from typing_extensions import Self
 
 import rmm
 
@@ -34,9 +26,6 @@
     from cudf.core.buffer.spill_manager import SpillManager
 
 
-T = TypeVar("T", bound="SpillableBuffer")
-
-
 def get_spillable_owner(data) -> Optional[SpillableBuffer]:
     """Get the spillable owner of `data`, if any exist
 
@@ -212,9 +201,7 @@ def _finalize_init(self, ptr_desc: Dict[str, Any], exposed: bool) -> None:
         self._manager.add(self)
 
     @classmethod
-    def _from_device_memory(
-        cls: Type[T], data: Any, *, exposed: bool = False
-    ) -> T:
+    def _from_device_memory(cls, data: Any, *, exposed: bool = False) -> Self:
         """Create a spillabe buffer from device memory.
 
         No data is being copied.
@@ -236,7 +223,7 @@ def _from_device_memory(
         return ret
 
     @classmethod
-    def _from_host_memory(cls: Type[T], data: Any) -> T:
+    def _from_host_memory(cls, data: Any) -> Self:
         """Create a spillabe buffer from host memory.
 
         Data must implement `__array_interface__`, the buffer protocol, and/or

python/cudf/cudf/core/column/categorical.py

@@ -11,6 +11,7 @@
 import pandas as pd
 import pyarrow as pa
 from numba import cuda
+from typing_extensions import Self
 
 import cudf
 from cudf import _lib as libcudf
@@ -716,7 +717,6 @@ def __init__(
         null_count: Optional[int] = None,
         children: Tuple["column.ColumnBase", ...] = (),
     ):
-
         if size is None:
             for child in children:
                 assert child.offset == 0
@@ -874,7 +874,7 @@ def _fill(
         begin: int,
         end: int,
         inplace: bool = False,
-    ) -> "column.ColumnBase":
+    ) -> Self:
         if end <= begin or begin >= self.size:
             return self if inplace else self.copy()
 
@@ -890,17 +890,20 @@ def _fill(
 
     def slice(
         self, start: int, stop: int, stride: Optional[int] = None
-    ) -> "column.ColumnBase":
+    ) -> Self:
         codes = self.codes.slice(start, stop, stride)
-        return cudf.core.column.build_categorical_column(
-            categories=self.categories,
-            codes=cudf.core.column.build_column(
-                codes.base_data, dtype=codes.dtype
+        return cast(
+            Self,
+            cudf.core.column.build_categorical_column(
+                categories=self.categories,
+                codes=cudf.core.column.build_column(
+                    codes.base_data, dtype=codes.dtype
+                ),
+                mask=codes.base_mask,
+                ordered=self.ordered,
+                size=codes.size,
+                offset=codes.offset,
             ),
-            mask=codes.base_mask,
-            ordered=self.ordered,
-            size=codes.size,
-            offset=codes.offset,
         )
 
     def _binaryop(self, other: ColumnBinaryOperand, op: str) -> ColumnBase:
@@ -1356,7 +1359,7 @@ def _get_decategorized_column(self) -> ColumnBase:
         out = out.set_mask(self.mask)
         return out
 
-    def copy(self, deep: bool = True) -> CategoricalColumn:
+    def copy(self, deep: bool = True) -> Self:
         result_col = super().copy(deep=deep)
         if deep:
             result_col.categories = libcudf.copying.copy_column(
@@ -1370,7 +1373,7 @@ def memory_usage(self) -> int:
 
     def _mimic_inplace(
         self, other_col: ColumnBase, inplace: bool = False
-    ) -> Optional[ColumnBase]:
+    ) -> Optional[Self]:
         out = super()._mimic_inplace(other_col, inplace=inplace)
         if inplace and isinstance(other_col, CategoricalColumn):
             self._codes = other_col._codes