REF: move registry, Registry to dtypes.base (#34830)

Author: jbrockmendel

pandas/api/extensions/__init__.py

@@ -4,7 +4,7 @@
 
 from pandas._libs.lib import no_default
 
-from pandas.core.dtypes.dtypes import ExtensionDtype, register_extension_dtype
+from pandas.core.dtypes.base import ExtensionDtype, register_extension_dtype
 
 from pandas.core.accessor import (
     register_dataframe_accessor,

pandas/core/arrays/integer.py

@@ -10,6 +10,7 @@
 from pandas.compat.numpy import function as nv
 from pandas.util._decorators import cache_readonly
 
+from pandas.core.dtypes.base import register_extension_dtype
 from pandas.core.dtypes.common import (
     is_bool_dtype,
     is_datetime64_dtype,
@@ -21,7 +22,6 @@
     is_object_dtype,
     pandas_dtype,
 )
-from pandas.core.dtypes.dtypes import register_extension_dtype
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import ops

pandas/core/arrays/sparse/dtype.py

@@ -9,7 +9,7 @@
 from pandas._typing import Dtype, DtypeObj
 from pandas.errors import PerformanceWarning
 
-from pandas.core.dtypes.base import ExtensionDtype
+from pandas.core.dtypes.base import ExtensionDtype, register_extension_dtype
 from pandas.core.dtypes.cast import astype_nansafe
 from pandas.core.dtypes.common import (
     is_bool_dtype,
@@ -19,7 +19,6 @@
     is_string_dtype,
     pandas_dtype,
 )
-from pandas.core.dtypes.dtypes import register_extension_dtype
 from pandas.core.dtypes.missing import isna, na_value_for_dtype
 
 if TYPE_CHECKING:

pandas/core/arrays/string_.py

@@ -5,9 +5,8 @@
 
 from pandas._libs import lib, missing as libmissing
 
-from pandas.core.dtypes.base import ExtensionDtype
+from pandas.core.dtypes.base import ExtensionDtype, register_extension_dtype
 from pandas.core.dtypes.common import pandas_dtype
-from pandas.core.dtypes.dtypes import register_extension_dtype
 from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries
 from pandas.core.dtypes.inference import is_array_like
 

pandas/core/construction.py

@@ -15,6 +15,7 @@
 from pandas._libs.tslibs import IncompatibleFrequency, OutOfBoundsDatetime
 from pandas._typing import AnyArrayLike, ArrayLike, Dtype, DtypeObj
 
+from pandas.core.dtypes.base import ExtensionDtype, registry
 from pandas.core.dtypes.cast import (
     construct_1d_arraylike_from_scalar,
     construct_1d_ndarray_preserving_na,
@@ -36,7 +37,6 @@
     is_object_dtype,
     is_timedelta64_ns_dtype,
 )
-from pandas.core.dtypes.dtypes import ExtensionDtype, registry
 from pandas.core.dtypes.generic import (
     ABCExtensionArray,
     ABCIndexClass,

pandas/core/dtypes/base.py

@@ -2,7 +2,7 @@
 Extend pandas with custom array types.
 """
 
-from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Type, Union
 
 import numpy as np
 
@@ -352,3 +352,92 @@ def _get_common_dtype(self, dtypes: List[DtypeObj]) -> Optional[DtypeObj]:
             return self
         else:
             return None
+
+
+def register_extension_dtype(cls: Type[ExtensionDtype]) -> Type[ExtensionDtype]:
+    """
+    Register an ExtensionType with pandas as class decorator.
+
+    .. versionadded:: 0.24.0
+
+    This enables operations like ``.astype(name)`` for the name
+    of the ExtensionDtype.
+
+    Returns
+    -------
+    callable
+        A class decorator.
+
+    Examples
+    --------
+    >>> from pandas.api.extensions import register_extension_dtype
+    >>> from pandas.api.extensions import ExtensionDtype
+    >>> @register_extension_dtype
+    ... class MyExtensionDtype(ExtensionDtype):
+    ...     name = "myextension"
+    """
+    registry.register(cls)
+    return cls
+
+
+class Registry:
+    """
+    Registry for dtype inference.
+
+    The registry allows one to map a string repr of a extension
+    dtype to an extension dtype. The string alias can be used in several
+    places, including
+
+    * Series and Index constructors
+    * :meth:`pandas.array`
+    * :meth:`pandas.Series.astype`
+
+    Multiple extension types can be registered.
+    These are tried in order.
+    """
+
+    def __init__(self):
+        self.dtypes: List[Type[ExtensionDtype]] = []
+
+    def register(self, dtype: Type[ExtensionDtype]) -> None:
+        """
+        Parameters
+        ----------
+        dtype : ExtensionDtype class
+        """
+        if not issubclass(dtype, ExtensionDtype):
+            raise ValueError("can only register pandas extension dtypes")
+
+        self.dtypes.append(dtype)
+
+    def find(
+        self, dtype: Union[Type[ExtensionDtype], str]
+    ) -> Optional[Type[ExtensionDtype]]:
+        """
+        Parameters
+        ----------
+        dtype : Type[ExtensionDtype] 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)
+            if issubclass(dtype_type, ExtensionDtype):
+                return dtype
+
+            return None
+
+        for dtype_type in self.dtypes:
+            try:
+                return dtype_type.construct_from_string(dtype)
+            except TypeError:
+                pass
+
+        return None
+
+
+registry = Registry()

pandas/core/dtypes/common.py

@@ -11,13 +11,13 @@
 from pandas._libs.tslibs import conversion
 from pandas._typing import ArrayLike, DtypeObj
 
+from pandas.core.dtypes.base import registry
 from pandas.core.dtypes.dtypes import (
     CategoricalDtype,
     DatetimeTZDtype,
     ExtensionDtype,
     IntervalDtype,
     PeriodDtype,
-    registry,
 )
 from pandas.core.dtypes.generic import ABCCategorical, ABCIndexClass
 from pandas.core.dtypes.inference import (  # noqa:F401

pandas/core/dtypes/dtypes.py

@@ -24,7 +24,7 @@
 from pandas._libs.tslibs.offsets import BaseOffset
 from pandas._typing import DtypeObj, Ordered
 
-from pandas.core.dtypes.base import ExtensionDtype
+from pandas.core.dtypes.base import ExtensionDtype, register_extension_dtype
 from pandas.core.dtypes.generic import ABCCategoricalIndex, ABCIndexClass
 from pandas.core.dtypes.inference import is_bool, is_list_like
 
@@ -40,95 +40,6 @@
 str_type = str
 
 
-def register_extension_dtype(cls: Type[ExtensionDtype]) -> Type[ExtensionDtype]:
-    """
-    Register an ExtensionType with pandas as class decorator.
-
-    .. versionadded:: 0.24.0
-
-    This enables operations like ``.astype(name)`` for the name
-    of the ExtensionDtype.
-
-    Returns
-    -------
-    callable
-        A class decorator.
-
-    Examples
-    --------
-    >>> from pandas.api.extensions import register_extension_dtype
-    >>> from pandas.api.extensions import ExtensionDtype
-    >>> @register_extension_dtype
-    ... class MyExtensionDtype(ExtensionDtype):
-    ...     pass
-    """
-    registry.register(cls)
-    return cls
-
-
-class Registry:
-    """
-    Registry for dtype inference.
-
-    The registry allows one to map a string repr of a extension
-    dtype to an extension dtype. The string alias can be used in several
-    places, including
-
-    * Series and Index constructors
-    * :meth:`pandas.array`
-    * :meth:`pandas.Series.astype`
-
-    Multiple extension types can be registered.
-    These are tried in order.
-    """
-
-    def __init__(self):
-        self.dtypes: List[Type[ExtensionDtype]] = []
-
-    def register(self, dtype: Type[ExtensionDtype]) -> None:
-        """
-        Parameters
-        ----------
-        dtype : ExtensionDtype class
-        """
-        if not issubclass(dtype, ExtensionDtype):
-            raise ValueError("can only register pandas extension dtypes")
-
-        self.dtypes.append(dtype)
-
-    def find(
-        self, dtype: Union[Type[ExtensionDtype], str]
-    ) -> Optional[Type[ExtensionDtype]]:
-        """
-        Parameters
-        ----------
-        dtype : Type[ExtensionDtype] 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)
-            if issubclass(dtype_type, ExtensionDtype):
-                return dtype
-
-            return None
-
-        for dtype_type in self.dtypes:
-            try:
-                return dtype_type.construct_from_string(dtype)
-            except TypeError:
-                pass
-
-        return None
-
-
-registry = Registry()
-
-
 class PandasExtensionDtype(ExtensionDtype):
     """
     A np.dtype duck-typed class, suitable for holding a custom dtype.

pandas/tests/arrays/test_array.py

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

pandas/tests/arrays/test_period.py

@@ -5,7 +5,8 @@
 from pandas._libs.tslibs.period import IncompatibleFrequency
 import pandas.util._test_decorators as td
 
-from pandas.core.dtypes.dtypes import PeriodDtype, registry
+from pandas.core.dtypes.base import registry
+from pandas.core.dtypes.dtypes import PeriodDtype
 
 import pandas as pd
 import pandas._testing as tm

pandas/tests/dtypes/test_dtypes.py

@@ -4,6 +4,7 @@
 import pytest
 import pytz
 
+from pandas.core.dtypes.base import registry
 from pandas.core.dtypes.common import (
     is_bool_dtype,
     is_categorical,
@@ -22,7 +23,6 @@
     DatetimeTZDtype,
     IntervalDtype,
     PeriodDtype,
-    registry,
 )
 
 import pandas as pd