Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
Jul 10, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pandas/api/extensions/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion pandas/core/arrays/integer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down
3 changes: 1 addition & 2 deletions pandas/core/arrays/sparse/dtype.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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:
Expand Down
3 changes: 1 addition & 2 deletions pandas/core/arrays/string_.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion pandas/core/construction.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down
91 changes: 90 additions & 1 deletion pandas/core/dtypes/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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]:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Perhaps even their own file?

Right now it might not merit its own file, but in the future I could see us having an entrypoints-based registry, which will add more lines.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah i agree, let's create registry.py

"""
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()
2 changes: 1 addition & 1 deletion pandas/core/dtypes/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
91 changes: 1 addition & 90 deletions pandas/core/dtypes/dtypes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion pandas/tests/arrays/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion pandas/tests/arrays/test_period.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion pandas/tests/dtypes/test_dtypes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -22,7 +23,6 @@
DatetimeTZDtype,
IntervalDtype,
PeriodDtype,
registry,
)

import pandas as pd
Expand Down