TYP: Fix some typehints for ExtensionDtype

Author: ozars

pandas/_typing.py

@@ -131,6 +131,9 @@
 # to maintain type information across generic functions and parametrization
 T = TypeVar("T")
 
+# To parameterize on same ExtensionDtype
+E = TypeVar("E", bound=ExtensionDtype)
+
 # used in decorators to preserve the signature of the function it decorates
 # see https://mypy.readthedocs.io/en/stable/generics.html#declaring-decorators
 FuncType = Callable[..., Any]

pandas/core/dtypes/base.py

@@ -13,6 +13,7 @@
 
 from pandas._typing import (
     DtypeObj,
+    E,
     type_t,
 )
 from pandas.errors import AbstractMethodError
@@ -151,7 +152,7 @@ def na_value(self) -> object:
         return np.nan
 
     @property
-    def type(self) -> type[Any]:
+    def type(self) -> type_t[Any]:
         """
         The scalar type for the array, e.g. ``int``
 
@@ -209,7 +210,7 @@ def construct_array_type(cls) -> type_t[ExtensionArray]:
         raise NotImplementedError
 
     @classmethod
-    def construct_from_string(cls, string: str):
+    def construct_from_string(cls, string: str) -> ExtensionDtype:
         r"""
         Construct this type from a string.
 
@@ -364,7 +365,7 @@ def _get_common_dtype(self, dtypes: list[DtypeObj]) -> DtypeObj | None:
             return None
 
 
-def register_extension_dtype(cls: type[ExtensionDtype]) -> type[ExtensionDtype]:
+def register_extension_dtype(cls: type[E]) -> type[E]:
     """
     Register an ExtensionType with pandas as class decorator.
 
@@ -420,7 +421,7 @@ def register(self, dtype: type[ExtensionDtype]) -> None:
 
         self.dtypes.append(dtype)
 
-    def find(self, dtype: type[ExtensionDtype] | str) -> type[ExtensionDtype] | None:
+    def find(self, dtype: type[E] | E | str) -> type[E] | E | ExtensionDtype | None:
         """
         Parameters
         ----------
@@ -431,10 +432,7 @@ def find(self, dtype: type[ExtensionDtype] | str) -> type[ExtensionDtype] | None
         return the first matching dtype, otherwise return None
         """
         if not isinstance(dtype, str):
-            dtype_type = dtype
-            if not isinstance(dtype, type):
-                dtype_type = type(dtype)
-            if issubclass(dtype_type, ExtensionDtype):
+            if isinstance(dtype, (ExtensionDtype, type(ExtensionDtype))):
                 return dtype
 
             return None