|
2 | 2 |
|
3 | 3 | import re
|
4 | 4 | import numpy as np
|
| 5 | +from collections import OrderedDict |
5 | 6 | from pandas import compat
|
6 | 7 | from pandas.core.dtypes.generic import ABCIndexClass, ABCCategoricalIndex
|
7 | 8 |
|
8 | 9 | from .base import ExtensionDtype, _DtypeOpsMixin
|
9 | 10 |
|
10 | 11 |
|
| 12 | +class Registry(object): |
| 13 | + """ Registry for dtype inference |
| 14 | +
|
| 15 | + We can directly construct dtypes in pandas_dtypes if they are |
| 16 | + a type; the registry allows us to register an extension dtype |
| 17 | + to try inference from a string or a dtype class |
| 18 | +
|
| 19 | + These are tried in order for inference. |
| 20 | + """ |
| 21 | + dtypes = OrderedDict() |
| 22 | + |
| 23 | + @classmethod |
| 24 | + def register(self, dtype, constructor=None): |
| 25 | + """ |
| 26 | + Parameters |
| 27 | + ---------- |
| 28 | + dtype : PandasExtension Dtype |
| 29 | + """ |
| 30 | + if not issubclass(dtype, (PandasExtensionDtype, ExtensionDtype)): |
| 31 | + raise ValueError("can only register pandas extension dtypes") |
| 32 | + |
| 33 | + if constructor is None: |
| 34 | + constructor = dtype.construct_from_string |
| 35 | + |
| 36 | + self.dtypes[dtype] = constructor |
| 37 | + |
| 38 | + def find(self, dtype): |
| 39 | + """ |
| 40 | + Parameters |
| 41 | + ---------- |
| 42 | + dtype : PandasExtensionDtype or string |
| 43 | +
|
| 44 | + Returns |
| 45 | + ------- |
| 46 | + return the first matching dtype, otherwise return None |
| 47 | + """ |
| 48 | + if not isinstance(dtype, compat.string_types): |
| 49 | + dtype_type = dtype |
| 50 | + if not isinstance(dtype, type): |
| 51 | + dtype_type = type(dtype) |
| 52 | + if issubclass(dtype_type, (PandasExtensionDtype, ExtensionDtype)): |
| 53 | + return dtype |
| 54 | + |
| 55 | + return None |
| 56 | + |
| 57 | + for dtype_type, constructor in self.dtypes.items(): |
| 58 | + try: |
| 59 | + return constructor(dtype) |
| 60 | + except TypeError: |
| 61 | + pass |
| 62 | + |
| 63 | + return None |
| 64 | + |
| 65 | + |
| 66 | +registry = Registry() |
| 67 | + |
| 68 | + |
11 | 69 | class PandasExtensionDtype(_DtypeOpsMixin):
|
12 | 70 | """
|
13 | 71 | A np.dtype duck-typed class, suitable for holding a custom dtype.
|
@@ -564,6 +622,17 @@ def construct_from_string(cls, string):
|
564 | 622 | pass
|
565 | 623 | raise TypeError("could not construct PeriodDtype")
|
566 | 624 |
|
| 625 | + @classmethod |
| 626 | + def construct_from_string_strict(cls, string): |
| 627 | + """ |
| 628 | + Strict construction from a string, raise a TypeError if not |
| 629 | + possible |
| 630 | + """ |
| 631 | + if string.startswith('period[') or string.startswith('Period['): |
| 632 | + # do not parse string like U as period[U] |
| 633 | + return PeriodDtype.construct_from_string(string) |
| 634 | + raise TypeError("could not construct PeriodDtype") |
| 635 | + |
567 | 636 | def __unicode__(self):
|
568 | 637 | return "period[{freq}]".format(freq=self.freq.freqstr)
|
569 | 638 |
|
@@ -683,6 +752,16 @@ def construct_from_string(cls, string):
|
683 | 752 | msg = "a string needs to be passed, got type {typ}"
|
684 | 753 | raise TypeError(msg.format(typ=type(string)))
|
685 | 754 |
|
| 755 | + @classmethod |
| 756 | + def construct_from_string_strict(cls, string): |
| 757 | + """ |
| 758 | + Strict construction from a string, raise a TypeError if not |
| 759 | + possible |
| 760 | + """ |
| 761 | + if string.startswith('interval') or string.startswith('Interval'): |
| 762 | + return IntervalDtype.construct_from_string(string) |
| 763 | + raise TypeError("cannot construct IntervalDtype") |
| 764 | + |
686 | 765 | def __unicode__(self):
|
687 | 766 | if self.subtype is None:
|
688 | 767 | return "interval"
|
@@ -723,3 +802,10 @@ def is_dtype(cls, dtype):
|
723 | 802 | else:
|
724 | 803 | return False
|
725 | 804 | return super(IntervalDtype, cls).is_dtype(dtype)
|
| 805 | + |
| 806 | + |
| 807 | +# register the dtypes in search order |
| 808 | +registry.register(DatetimeTZDtype) |
| 809 | +registry.register(PeriodDtype, PeriodDtype.construct_from_string_strict) |
| 810 | +registry.register(IntervalDtype, IntervalDtype.construct_from_string_strict) |
| 811 | +registry.register(CategoricalDtype) |
0 commit comments