lose less ExtensionArrayT. Make registry private. consolidate overload

Author: Dr-Irv

pandas/core/arrays/_mixins.py

@@ -255,13 +255,7 @@ def __getitem__(self: NDArrayBackedExtensionArrayT, key: int) -> Any:
 
     @overload
     def __getitem__(
-        self: NDArrayBackedExtensionArrayT, key: slice
-    ) -> NDArrayBackedExtensionArrayT:
-        ...
-
-    @overload
-    def __getitem__(
-        self: NDArrayBackedExtensionArrayT, key: np.ndarray
+        self: NDArrayBackedExtensionArrayT, key: Union[slice, np.ndarray]
     ) -> NDArrayBackedExtensionArrayT:
         ...
 

pandas/core/arrays/base.py

@@ -289,19 +289,15 @@ def _from_factorized(cls, values, original):
     # ------------------------------------------------------------------------
 
     @overload
-    def __getitem__(self: ExtensionArrayT, item: int) -> Any:
+    def __getitem__(self, item: int) -> Any:
         ...
 
     @overload
-    def __getitem__(self: ExtensionArrayT, item: slice) -> ExtensionArrayT:
-        ...
-
-    @overload
-    def __getitem__(self: ExtensionArrayT, item: np.ndarray) -> ExtensionArrayT:
+    def __getitem__(self, item: Union[slice, np.ndarray]) -> ExtensionArray:
         ...
 
     def __getitem__(
-        self: ExtensionArrayT, item: Union[int, slice, np.ndarray]
+        self, item: Union[int, slice, np.ndarray]
     ) -> Union[ExtensionArray, Any]:
         """
         Select a subset of self.
@@ -686,11 +682,11 @@ def argmax(self, skipna: bool = True) -> int:
         return nargminmax(self, "argmax")
 
     def fillna(
-        self: ExtensionArrayT,
+        self: ExtensionArray,
         value: Optional[Union[Any, ArrayLike]] = None,
         method: Optional[Literal["backfill", "bfill", "ffill", "pad"]] = None,
         limit: Optional[int] = None,
-    ) -> ExtensionArrayT:
+    ) -> ExtensionArray:
         """
         Fill NA/NaN values using the specified method.
 
@@ -746,8 +742,8 @@ def dropna(self) -> ExtensionArrayT:
         return self[~self.isna()]
 
     def shift(
-        self: ExtensionArrayT, periods: int = 1, fill_value: object = None
-    ) -> ExtensionArrayT:
+        self: ExtensionArray, periods: int = 1, fill_value: object = None
+    ) -> ExtensionArray:
         """
         Shift values by desired number.
 
@@ -795,13 +791,13 @@ def shift(
         )
         if periods > 0:
             a = empty
-            b: ExtensionArrayT = self[:-periods]
+            b = self[:-periods]
         else:
             a = self[abs(periods) :]
             b = empty
         return self._concat_same_type([a, b])
 
-    def unique(self: ExtensionArrayT) -> ExtensionArrayT:
+    def unique(self: ExtensionArray) -> ExtensionArray:
         """
         Compute the ExtensionArray of unique values.
 
@@ -815,7 +811,7 @@ def unique(self: ExtensionArrayT) -> ExtensionArrayT:
     def searchsorted(
         self,
         value: ArrayLike,
-        side: Optional[Literal["left", "right"]] = "left",
+        side: Literal["left", "right"] = "left",
         sorter: Optional[ArrayLike] = None,
     ) -> np.ndarray:
         """
@@ -1033,10 +1029,10 @@ def factorize(self, na_sentinel: int = -1) -> Tuple[np.ndarray, ExtensionArray]:
     @Substitution(klass="ExtensionArray")
     @Appender(_extension_array_shared_docs["repeat"])
     def repeat(
-        self: ExtensionArrayT,
+        self: ExtensionArray,
         repeats: Union[int, Sequence[int]],
         axis: Literal[None] = None,
-    ) -> ExtensionArrayT:
+    ) -> ExtensionArray:
         nv.validate_repeat((), {"axis": axis})
         ind = np.arange(len(self)).repeat(repeats)
         return self.take(ind)
@@ -1046,12 +1042,12 @@ def repeat(
     # ------------------------------------------------------------------------
 
     def take(
-        self: ExtensionArrayT,
+        self: ExtensionArray,
         indices: Sequence[int],
         *,
         allow_fill: bool = False,
         fill_value: Any = None,
-    ) -> ExtensionArrayT:
+    ) -> ExtensionArray:
         """
         Take elements from an array.
 
@@ -1140,7 +1136,7 @@ def take(self, indices, allow_fill=False, fill_value=None):
         # pandas.api.extensions.take
         raise AbstractMethodError(self)
 
-    def copy(self: ExtensionArrayT) -> ExtensionArrayT:
+    def copy(self: ExtensionArray) -> ExtensionArray:
         """
         Return a copy of the array.
 
@@ -1220,7 +1216,7 @@ def _formatter(self, boxed: bool = False) -> Callable[[Any], Optional[str]]:
     # Reshaping
     # ------------------------------------------------------------------------
 
-    def transpose(self: ExtensionArrayT, *axes: int) -> ExtensionArrayT:
+    def transpose(self: ExtensionArray, *axes: int) -> ExtensionArray:
         """
         Return a transposed view on this array.
 
@@ -1230,12 +1226,12 @@ def transpose(self: ExtensionArrayT, *axes: int) -> ExtensionArrayT:
         return self[:]
 
     @property
-    def T(self: ExtensionArrayT) -> ExtensionArrayT:
+    def T(self: ExtensionArray) -> ExtensionArray:
         return self.transpose()
 
     def ravel(
-        self: ExtensionArrayT, order: Optional[Literal["C", "F", "A", "K"]] = "C"
-    ) -> ExtensionArrayT:
+        self: ExtensionArray, order: Optional[Literal["C", "F", "A", "K"]] = "C"
+    ) -> ExtensionArray:
         """
         Return a flattened view on this array.
 
@@ -1257,7 +1253,7 @@ def ravel(
     @classmethod
     def _concat_same_type(
         cls: Type[ExtensionArrayT], to_concat: Sequence[ExtensionArrayT]
-    ):
+    ) -> ExtensionArrayT:
         """
         Concatenate multiple array of this dtype.
 
@@ -1315,9 +1311,7 @@ def __hash__(self) -> int:
     # ------------------------------------------------------------------------
     # Non-Optimized Default Methods
 
-    def delete(
-        self: ExtensionArrayT, loc: Union[int, Sequence[int]]
-    ) -> ExtensionArrayT:
+    def delete(self: ExtensionArray, loc: Union[int, Sequence[int]]) -> ExtensionArray:
         indexer = np.delete(np.arange(len(self)), loc)
         return self.take(indexer)
 

pandas/core/construction.py

@@ -33,7 +33,7 @@
 
 from pandas.core.dtypes.base import (
     ExtensionDtype,
-    registry,
+    _registry as registry,
 )
 from pandas.core.dtypes.cast import (
     construct_1d_arraylike_from_scalar,

pandas/core/dtypes/base.py

@@ -389,7 +389,7 @@ def register_extension_dtype(cls: Type[ExtensionDtype]) -> Type[ExtensionDtype]:
     ... class MyExtensionDtype(ExtensionDtype):
     ...     name = "myextension"
     """
-    registry.register(cls)
+    _registry.register(cls)
     return cls
 
 
@@ -451,4 +451,4 @@ def find(self, dtype: Union[Type[ExtensionDtype], str]) -> Optional[ExtensionDty
         return None
 
 
-registry: Registry = Registry()
+_registry = Registry()

pandas/core/dtypes/common.py

@@ -23,7 +23,7 @@
     Optional,
 )
 
-from pandas.core.dtypes.base import registry
+from pandas.core.dtypes.base import _registry as registry
 from pandas.core.dtypes.dtypes import (
     CategoricalDtype,
     DatetimeTZDtype,

pandas/tests/arrays/test_array.py

@@ -5,7 +5,7 @@
 import pytest
 import pytz
 
-from pandas.core.dtypes.base import registry
+from pandas.core.dtypes.base import _registry as registry
 
 import pandas as pd
 import pandas._testing as tm

pandas/tests/arrays/test_period.py

@@ -4,7 +4,7 @@
 from pandas._libs.tslibs import iNaT
 from pandas._libs.tslibs.period import IncompatibleFrequency
 
-from pandas.core.dtypes.base import registry
+from pandas.core.dtypes.base import _registry as registry
 from pandas.core.dtypes.dtypes import PeriodDtype
 
 import pandas as pd

pandas/tests/dtypes/test_dtypes.py

@@ -4,7 +4,7 @@
 import pytest
 import pytz
 
-from pandas.core.dtypes.base import registry
+from pandas.core.dtypes.base import _registry as registry
 from pandas.core.dtypes.common import (
     is_bool_dtype,
     is_categorical,

pandas/tests/frame/indexing/test_setitem.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-from pandas.core.dtypes.base import registry as ea_registry
+from pandas.core.dtypes.base import _registry as ea_registry
 from pandas.core.dtypes.common import (
     is_categorical_dtype,
     is_interval_dtype,