TYP: Fix typing in ExtensionDtype registry

pandas/core/dtypes/base.py

@@ -8,6 +8,7 @@
     TYPE_CHECKING,
     Any,
     TypeVar,
+    overload,
 )
 
 import numpy as np
@@ -28,7 +29,7 @@
     from pandas.core.arrays import ExtensionArray
 
     # To parameterize on same ExtensionDtype
-    E = TypeVar("E", bound="ExtensionDtype")
+    ExtensionDtypeT = TypeVar("ExtensionDtypeT", bound="ExtensionDtype")
 
 
 class ExtensionDtype:
@@ -202,7 +203,7 @@ def names(self) -> list[str] | None:
         return None
 
     @classmethod
-    def construct_array_type(cls) -> type_t[ExtensionArray]:
+    def construct_array_type(cls: type_t[ExtensionDtypeT]) -> type_t[ExtensionArray]:
         """
         Return the array type associated with this dtype.
 
@@ -213,7 +214,9 @@ def construct_array_type(cls) -> type_t[ExtensionArray]:
         raise NotImplementedError
 
     @classmethod
-    def construct_from_string(cls, string: str):
+    def construct_from_string(
+        cls: type_t[ExtensionDtypeT], string: str
+    ) -> ExtensionDtype:
         r"""
         Construct this type from a string.
 
@@ -268,7 +271,7 @@ def construct_from_string(cls, string: str):
         return cls()
 
     @classmethod
-    def is_dtype(cls, dtype: object) -> bool:
+    def is_dtype(cls: type_t[ExtensionDtypeT], dtype: object) -> bool:
         """
         Check if we match 'dtype'.
 
@@ -368,7 +371,7 @@ def _get_common_dtype(self, dtypes: list[DtypeObj]) -> DtypeObj | None:
             return None
 
 
-def register_extension_dtype(cls: type[E]) -> type[E]:
+def register_extension_dtype(cls: type[ExtensionDtypeT]) -> type[ExtensionDtypeT]:
     """
     Register an ExtensionType with pandas as class decorator.
 
@@ -424,20 +427,35 @@ def register(self, dtype: type[ExtensionDtype]) -> None:
 
         self.dtypes.append(dtype)
 
-    def find(self, dtype: type[ExtensionDtype] | str) -> type[ExtensionDtype] | None:
+    @overload
+    def find(self, dtype: type[ExtensionDtype]) -> ExtensionDtype:
+        ...
+
+    @overload
+    def find(self, dtype: ExtensionDtype) -> ExtensionDtype:
+        ...
+
+    @overload
+    def find(self, dtype: str) -> ExtensionDtype | None:
+        ...
+
+    def find(
+        self, dtype: type[ExtensionDtype] | ExtensionDtype | str
+    ) -> type[ExtensionDtype] | ExtensionDtype | None:
         """
         Parameters
         ----------
-        dtype : Type[ExtensionDtype] or str
+        dtype : ExtensionDtype class or instance or str
 
         Returns
         -------
         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)
+            else:
+                dtype_type = dtype
             if issubclass(dtype_type, ExtensionDtype):
                 return dtype
 

pandas/core/dtypes/common.py

@@ -1789,9 +1789,7 @@ def pandas_dtype(dtype) -> DtypeObj:
     # registered extension types
     result = registry.find(dtype)
     if result is not None:
-        # error: Incompatible return value type (got "Type[ExtensionDtype]",
-        # expected "Union[dtype, ExtensionDtype]")
-        return result  # type: ignore[return-value]
+        return result
 
     # try a numpy dtype
     # raise a consistent TypeError if failed