|
2 | 2 | Extend pandas with custom array types.
|
3 | 3 | """
|
4 | 4 |
|
5 |
| -from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Type |
| 5 | +from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Type, Union |
6 | 6 |
|
7 | 7 | import numpy as np
|
8 | 8 |
|
@@ -352,3 +352,92 @@ def _get_common_dtype(self, dtypes: List[DtypeObj]) -> Optional[DtypeObj]:
|
352 | 352 | return self
|
353 | 353 | else:
|
354 | 354 | 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() |
0 commit comments