Skip to content

Commit 4fc4622

Browse files
authored
REF: move registry, Registry to dtypes.base (#34830)
1 parent 8fd4cd3 commit 4fc4622

File tree

11 files changed

+101
-102
lines changed

11 files changed

+101
-102
lines changed

pandas/api/extensions/__init__.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
from pandas._libs.lib import no_default
66

7-
from pandas.core.dtypes.dtypes import ExtensionDtype, register_extension_dtype
7+
from pandas.core.dtypes.base import ExtensionDtype, register_extension_dtype
88

99
from pandas.core.accessor import (
1010
register_dataframe_accessor,

pandas/core/arrays/integer.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from pandas.compat.numpy import function as nv
1111
from pandas.util._decorators import cache_readonly
1212

13+
from pandas.core.dtypes.base import register_extension_dtype
1314
from pandas.core.dtypes.common import (
1415
is_bool_dtype,
1516
is_datetime64_dtype,
@@ -21,7 +22,6 @@
2122
is_object_dtype,
2223
pandas_dtype,
2324
)
24-
from pandas.core.dtypes.dtypes import register_extension_dtype
2525
from pandas.core.dtypes.missing import isna
2626

2727
from pandas.core import ops

pandas/core/arrays/sparse/dtype.py

+1-2
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from pandas._typing import Dtype, DtypeObj
1010
from pandas.errors import PerformanceWarning
1111

12-
from pandas.core.dtypes.base import ExtensionDtype
12+
from pandas.core.dtypes.base import ExtensionDtype, register_extension_dtype
1313
from pandas.core.dtypes.cast import astype_nansafe
1414
from pandas.core.dtypes.common import (
1515
is_bool_dtype,
@@ -19,7 +19,6 @@
1919
is_string_dtype,
2020
pandas_dtype,
2121
)
22-
from pandas.core.dtypes.dtypes import register_extension_dtype
2322
from pandas.core.dtypes.missing import isna, na_value_for_dtype
2423

2524
if TYPE_CHECKING:

pandas/core/arrays/string_.py

+1-2
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,8 @@
55

66
from pandas._libs import lib, missing as libmissing
77

8-
from pandas.core.dtypes.base import ExtensionDtype
8+
from pandas.core.dtypes.base import ExtensionDtype, register_extension_dtype
99
from pandas.core.dtypes.common import pandas_dtype
10-
from pandas.core.dtypes.dtypes import register_extension_dtype
1110
from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries
1211
from pandas.core.dtypes.inference import is_array_like
1312

pandas/core/construction.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
from pandas._libs.tslibs import IncompatibleFrequency, OutOfBoundsDatetime
1616
from pandas._typing import AnyArrayLike, ArrayLike, Dtype, DtypeObj
1717

18+
from pandas.core.dtypes.base import ExtensionDtype, registry
1819
from pandas.core.dtypes.cast import (
1920
construct_1d_arraylike_from_scalar,
2021
construct_1d_ndarray_preserving_na,
@@ -36,7 +37,6 @@
3637
is_object_dtype,
3738
is_timedelta64_ns_dtype,
3839
)
39-
from pandas.core.dtypes.dtypes import ExtensionDtype, registry
4040
from pandas.core.dtypes.generic import (
4141
ABCExtensionArray,
4242
ABCIndexClass,

pandas/core/dtypes/base.py

+90-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
Extend pandas with custom array types.
33
"""
44

5-
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Type
5+
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Type, Union
66

77
import numpy as np
88

@@ -352,3 +352,92 @@ def _get_common_dtype(self, dtypes: List[DtypeObj]) -> Optional[DtypeObj]:
352352
return self
353353
else:
354354
return None
355+
356+
357+
def register_extension_dtype(cls: Type[ExtensionDtype]) -> Type[ExtensionDtype]:
358+
"""
359+
Register an ExtensionType with pandas as class decorator.
360+
361+
.. versionadded:: 0.24.0
362+
363+
This enables operations like ``.astype(name)`` for the name
364+
of the ExtensionDtype.
365+
366+
Returns
367+
-------
368+
callable
369+
A class decorator.
370+
371+
Examples
372+
--------
373+
>>> from pandas.api.extensions import register_extension_dtype
374+
>>> from pandas.api.extensions import ExtensionDtype
375+
>>> @register_extension_dtype
376+
... class MyExtensionDtype(ExtensionDtype):
377+
... name = "myextension"
378+
"""
379+
registry.register(cls)
380+
return cls
381+
382+
383+
class Registry:
384+
"""
385+
Registry for dtype inference.
386+
387+
The registry allows one to map a string repr of a extension
388+
dtype to an extension dtype. The string alias can be used in several
389+
places, including
390+
391+
* Series and Index constructors
392+
* :meth:`pandas.array`
393+
* :meth:`pandas.Series.astype`
394+
395+
Multiple extension types can be registered.
396+
These are tried in order.
397+
"""
398+
399+
def __init__(self):
400+
self.dtypes: List[Type[ExtensionDtype]] = []
401+
402+
def register(self, dtype: Type[ExtensionDtype]) -> None:
403+
"""
404+
Parameters
405+
----------
406+
dtype : ExtensionDtype class
407+
"""
408+
if not issubclass(dtype, ExtensionDtype):
409+
raise ValueError("can only register pandas extension dtypes")
410+
411+
self.dtypes.append(dtype)
412+
413+
def find(
414+
self, dtype: Union[Type[ExtensionDtype], str]
415+
) -> Optional[Type[ExtensionDtype]]:
416+
"""
417+
Parameters
418+
----------
419+
dtype : Type[ExtensionDtype] or str
420+
421+
Returns
422+
-------
423+
return the first matching dtype, otherwise return None
424+
"""
425+
if not isinstance(dtype, str):
426+
dtype_type = dtype
427+
if not isinstance(dtype, type):
428+
dtype_type = type(dtype)
429+
if issubclass(dtype_type, ExtensionDtype):
430+
return dtype
431+
432+
return None
433+
434+
for dtype_type in self.dtypes:
435+
try:
436+
return dtype_type.construct_from_string(dtype)
437+
except TypeError:
438+
pass
439+
440+
return None
441+
442+
443+
registry = Registry()

pandas/core/dtypes/common.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -11,13 +11,13 @@
1111
from pandas._libs.tslibs import conversion
1212
from pandas._typing import ArrayLike, DtypeObj
1313

14+
from pandas.core.dtypes.base import registry
1415
from pandas.core.dtypes.dtypes import (
1516
CategoricalDtype,
1617
DatetimeTZDtype,
1718
ExtensionDtype,
1819
IntervalDtype,
1920
PeriodDtype,
20-
registry,
2121
)
2222
from pandas.core.dtypes.generic import ABCCategorical, ABCIndexClass
2323
from pandas.core.dtypes.inference import ( # noqa:F401

pandas/core/dtypes/dtypes.py

+1-90
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
from pandas._libs.tslibs.offsets import BaseOffset
2525
from pandas._typing import DtypeObj, Ordered
2626

27-
from pandas.core.dtypes.base import ExtensionDtype
27+
from pandas.core.dtypes.base import ExtensionDtype, register_extension_dtype
2828
from pandas.core.dtypes.generic import ABCCategoricalIndex, ABCIndexClass
2929
from pandas.core.dtypes.inference import is_bool, is_list_like
3030

@@ -40,95 +40,6 @@
4040
str_type = str
4141

4242

43-
def register_extension_dtype(cls: Type[ExtensionDtype]) -> Type[ExtensionDtype]:
44-
"""
45-
Register an ExtensionType with pandas as class decorator.
46-
47-
.. versionadded:: 0.24.0
48-
49-
This enables operations like ``.astype(name)`` for the name
50-
of the ExtensionDtype.
51-
52-
Returns
53-
-------
54-
callable
55-
A class decorator.
56-
57-
Examples
58-
--------
59-
>>> from pandas.api.extensions import register_extension_dtype
60-
>>> from pandas.api.extensions import ExtensionDtype
61-
>>> @register_extension_dtype
62-
... class MyExtensionDtype(ExtensionDtype):
63-
... pass
64-
"""
65-
registry.register(cls)
66-
return cls
67-
68-
69-
class Registry:
70-
"""
71-
Registry for dtype inference.
72-
73-
The registry allows one to map a string repr of a extension
74-
dtype to an extension dtype. The string alias can be used in several
75-
places, including
76-
77-
* Series and Index constructors
78-
* :meth:`pandas.array`
79-
* :meth:`pandas.Series.astype`
80-
81-
Multiple extension types can be registered.
82-
These are tried in order.
83-
"""
84-
85-
def __init__(self):
86-
self.dtypes: List[Type[ExtensionDtype]] = []
87-
88-
def register(self, dtype: Type[ExtensionDtype]) -> None:
89-
"""
90-
Parameters
91-
----------
92-
dtype : ExtensionDtype class
93-
"""
94-
if not issubclass(dtype, ExtensionDtype):
95-
raise ValueError("can only register pandas extension dtypes")
96-
97-
self.dtypes.append(dtype)
98-
99-
def find(
100-
self, dtype: Union[Type[ExtensionDtype], str]
101-
) -> Optional[Type[ExtensionDtype]]:
102-
"""
103-
Parameters
104-
----------
105-
dtype : Type[ExtensionDtype] or str
106-
107-
Returns
108-
-------
109-
return the first matching dtype, otherwise return None
110-
"""
111-
if not isinstance(dtype, str):
112-
dtype_type = dtype
113-
if not isinstance(dtype, type):
114-
dtype_type = type(dtype)
115-
if issubclass(dtype_type, ExtensionDtype):
116-
return dtype
117-
118-
return None
119-
120-
for dtype_type in self.dtypes:
121-
try:
122-
return dtype_type.construct_from_string(dtype)
123-
except TypeError:
124-
pass
125-
126-
return None
127-
128-
129-
registry = Registry()
130-
131-
13243
class PandasExtensionDtype(ExtensionDtype):
13344
"""
13445
A np.dtype duck-typed class, suitable for holding a custom dtype.

pandas/tests/arrays/test_array.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import pytest
66
import pytz
77

8-
from pandas.core.dtypes.dtypes import registry
8+
from pandas.core.dtypes.base import registry
99

1010
import pandas as pd
1111
import pandas._testing as tm

pandas/tests/arrays/test_period.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@
55
from pandas._libs.tslibs.period import IncompatibleFrequency
66
import pandas.util._test_decorators as td
77

8-
from pandas.core.dtypes.dtypes import PeriodDtype, registry
8+
from pandas.core.dtypes.base import registry
9+
from pandas.core.dtypes.dtypes import PeriodDtype
910

1011
import pandas as pd
1112
import pandas._testing as tm

pandas/tests/dtypes/test_dtypes.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import pytest
55
import pytz
66

7+
from pandas.core.dtypes.base import registry
78
from pandas.core.dtypes.common import (
89
is_bool_dtype,
910
is_categorical,
@@ -22,7 +23,6 @@
2223
DatetimeTZDtype,
2324
IntervalDtype,
2425
PeriodDtype,
25-
registry,
2626
)
2727

2828
import pandas as pd

0 commit comments

Comments
 (0)