diff --git a/pandas/api/extensions/__init__.py b/pandas/api/extensions/__init__.py index 3019dd0e9b371..401e7081d2422 100644 --- a/pandas/api/extensions/__init__.py +++ b/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, diff --git a/pandas/core/arrays/integer.py b/pandas/core/arrays/integer.py index b5cb681812939..b0958af41158c 100644 --- a/pandas/core/arrays/integer.py +++ b/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 diff --git a/pandas/core/arrays/sparse/dtype.py b/pandas/core/arrays/sparse/dtype.py index b3da9cbeb44af..ccf2825162f51 100644 --- a/pandas/core/arrays/sparse/dtype.py +++ b/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: diff --git a/pandas/core/arrays/string_.py b/pandas/core/arrays/string_.py index ac501a8afbe09..5104e3f12f5b4 100644 --- a/pandas/core/arrays/string_.py +++ b/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 diff --git a/pandas/core/construction.py b/pandas/core/construction.py index 9ac661f97a56e..6c58698989e96 100644 --- a/pandas/core/construction.py +++ b/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, diff --git a/pandas/core/dtypes/base.py b/pandas/core/dtypes/base.py index 2d81dd4d884a3..07c73876954d0 100644 --- a/pandas/core/dtypes/base.py +++ b/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() diff --git a/pandas/core/dtypes/common.py b/pandas/core/dtypes/common.py index 9e960375e9bf4..a2ca4d84b2bf6 100644 --- a/pandas/core/dtypes/common.py +++ b/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 diff --git a/pandas/core/dtypes/dtypes.py b/pandas/core/dtypes/dtypes.py index a9d2430717e4f..22480fbc47508 100644 --- a/pandas/core/dtypes/dtypes.py +++ b/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. diff --git a/pandas/tests/arrays/test_array.py b/pandas/tests/arrays/test_array.py index ad6e6e4a98057..a0525aa511ee2 100644 --- a/pandas/tests/arrays/test_array.py +++ b/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 diff --git a/pandas/tests/arrays/test_period.py b/pandas/tests/arrays/test_period.py index 27e6334788284..8887dd0278afe 100644 --- a/pandas/tests/arrays/test_period.py +++ b/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 diff --git a/pandas/tests/dtypes/test_dtypes.py b/pandas/tests/dtypes/test_dtypes.py index b1fe673e9e2f1..a58dc5e5ec74a 100644 --- a/pandas/tests/dtypes/test_dtypes.py +++ b/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