|
| 1 | +import datetime as dt |
| 2 | +from typing import ( |
| 3 | + Any, |
| 4 | + Optional, |
| 5 | + Sequence, |
| 6 | + Tuple, |
| 7 | + Union, |
| 8 | + cast, |
| 9 | +) |
| 10 | + |
| 11 | +import numpy as np |
| 12 | + |
| 13 | +from pandas._typing import ( |
| 14 | + Dtype, |
| 15 | + PositionalIndexer, |
| 16 | +) |
| 17 | + |
| 18 | +from pandas.core.dtypes.dtypes import register_extension_dtype |
| 19 | + |
| 20 | +from pandas.api.extensions import ( |
| 21 | + ExtensionArray, |
| 22 | + ExtensionDtype, |
| 23 | +) |
| 24 | +from pandas.api.types import pandas_dtype |
| 25 | + |
| 26 | + |
| 27 | +@register_extension_dtype |
| 28 | +class DateDtype(ExtensionDtype): |
| 29 | + @property |
| 30 | + def type(self): |
| 31 | + return dt.date |
| 32 | + |
| 33 | + @property |
| 34 | + def name(self): |
| 35 | + return "DateDtype" |
| 36 | + |
| 37 | + @classmethod |
| 38 | + def construct_from_string(cls, string: str): |
| 39 | + if not isinstance(string, str): |
| 40 | + raise TypeError( |
| 41 | + f"'construct_from_string' expects a string, got {type(string)}" |
| 42 | + ) |
| 43 | + |
| 44 | + if string == cls.__name__: |
| 45 | + return cls() |
| 46 | + else: |
| 47 | + raise TypeError(f"Cannot construct a '{cls.__name__}' from '{string}'") |
| 48 | + |
| 49 | + @classmethod |
| 50 | + def construct_array_type(cls): |
| 51 | + return DateArray |
| 52 | + |
| 53 | + @property |
| 54 | + def na_value(self): |
| 55 | + return dt.date.min |
| 56 | + |
| 57 | + def __repr__(self) -> str: |
| 58 | + return self.name |
| 59 | + |
| 60 | + |
| 61 | +class DateArray(ExtensionArray): |
| 62 | + def __init__( |
| 63 | + self, |
| 64 | + dates: Union[ |
| 65 | + dt.date, |
| 66 | + Sequence[dt.date], |
| 67 | + Tuple[np.ndarray, np.ndarray, np.ndarray], |
| 68 | + np.ndarray, |
| 69 | + ], |
| 70 | + ) -> None: |
| 71 | + if isinstance(dates, dt.date): |
| 72 | + self._year = np.array([dates.year]) |
| 73 | + self._month = np.array([dates.month]) |
| 74 | + self._day = np.array([dates.year]) |
| 75 | + return |
| 76 | + |
| 77 | + ldates = len(dates) |
| 78 | + if isinstance(dates, list): |
| 79 | + # pre-allocate the arrays since we know the size before hand |
| 80 | + self._year = np.zeros(ldates, dtype=np.uint16) # 65535 (0, 9999) |
| 81 | + self._month = np.zeros(ldates, dtype=np.uint8) # 255 (1, 31) |
| 82 | + self._day = np.zeros(ldates, dtype=np.uint8) # 255 (1, 12) |
| 83 | + # populate them |
| 84 | + for i, (y, m, d) in enumerate( |
| 85 | + map(lambda date: (date.year, date.month, date.day), dates) |
| 86 | + ): |
| 87 | + self._year[i] = y |
| 88 | + self._month[i] = m |
| 89 | + self._day[i] = d |
| 90 | + |
| 91 | + elif isinstance(dates, tuple): |
| 92 | + # only support triples |
| 93 | + if ldates != 3: |
| 94 | + raise ValueError("only triples are valid") |
| 95 | + # check if all elements have the same type |
| 96 | + if any(map(lambda x: not isinstance(x, np.ndarray), dates)): |
| 97 | + raise TypeError("invalid type") |
| 98 | + ly, lm, ld = (len(cast(np.ndarray, d)) for d in dates) |
| 99 | + if not ly == lm == ld: |
| 100 | + raise ValueError( |
| 101 | + f"tuple members must have the same length: {(ly, lm, ld)}" |
| 102 | + ) |
| 103 | + self._year = dates[0].astype(np.uint16) |
| 104 | + self._month = dates[1].astype(np.uint8) |
| 105 | + self._day = dates[2].astype(np.uint8) |
| 106 | + |
| 107 | + elif isinstance(dates, np.ndarray) and dates.dtype == "U10": |
| 108 | + self._year = np.zeros(ldates, dtype=np.uint16) # 65535 (0, 9999) |
| 109 | + self._month = np.zeros(ldates, dtype=np.uint8) # 255 (1, 31) |
| 110 | + self._day = np.zeros(ldates, dtype=np.uint8) # 255 (1, 12) |
| 111 | + |
| 112 | + for (i,), (y, m, d) in np.ndenumerate(np.char.split(dates, sep="-")): |
| 113 | + self._year[i] = int(y) |
| 114 | + self._month[i] = int(m) |
| 115 | + self._day[i] = int(d) |
| 116 | + |
| 117 | + else: |
| 118 | + raise TypeError(f"{type(dates)} is not supported") |
| 119 | + |
| 120 | + @property |
| 121 | + def dtype(self) -> ExtensionDtype: |
| 122 | + return DateDtype() |
| 123 | + |
| 124 | + def astype(self, dtype, copy=True): |
| 125 | + dtype = pandas_dtype(dtype) |
| 126 | + |
| 127 | + if isinstance(dtype, DateDtype): |
| 128 | + data = self.copy() if copy else self |
| 129 | + else: |
| 130 | + data = self.to_numpy(dtype=dtype, copy=copy, na_value=dt.date.min) |
| 131 | + |
| 132 | + return data |
| 133 | + |
| 134 | + @property |
| 135 | + def nbytes(self) -> int: |
| 136 | + return self._year.nbytes + self._month.nbytes + self._day.nbytes |
| 137 | + |
| 138 | + def __len__(self) -> int: |
| 139 | + return len(self._year) # all 3 arrays are enforced to have the same length |
| 140 | + |
| 141 | + def __getitem__(self, item: PositionalIndexer): |
| 142 | + if isinstance(item, int): |
| 143 | + return dt.date(self._year[item], self._month[item], self._day[item]) |
| 144 | + else: |
| 145 | + raise NotImplementedError("only ints are supported as indexes") |
| 146 | + |
| 147 | + def __setitem__(self, key: Union[int, slice, np.ndarray], value: Any): |
| 148 | + if not isinstance(key, int): |
| 149 | + raise NotImplementedError("only ints are supported as indexes") |
| 150 | + |
| 151 | + if not isinstance(value, dt.date): |
| 152 | + raise TypeError("you can only set datetime.date types") |
| 153 | + |
| 154 | + self._year[key] = value.year |
| 155 | + self._month[key] = value.month |
| 156 | + self._day[key] = value.day |
| 157 | + |
| 158 | + def __repr__(self) -> str: |
| 159 | + return f"DateArray{list(zip(self._year, self._month, self._day))}" |
| 160 | + |
| 161 | + def copy(self) -> "DateArray": |
| 162 | + return DateArray((self._year.copy(), self._month.copy(), self._day.copy())) |
| 163 | + |
| 164 | + def isna(self) -> np.ndarray: |
| 165 | + return np.logical_and( |
| 166 | + np.logical_and( |
| 167 | + self._year == dt.date.min.year, self._month == dt.date.min.month |
| 168 | + ), |
| 169 | + self._day == dt.date.min.day, |
| 170 | + ) |
| 171 | + |
| 172 | + @classmethod |
| 173 | + def _from_sequence(cls, scalars, *, dtype: Optional[Dtype] = None, copy=False): |
| 174 | + if isinstance(scalars, dt.date): |
| 175 | + pass |
| 176 | + elif isinstance(scalars, DateArray): |
| 177 | + pass |
| 178 | + elif isinstance(scalars, np.ndarray): |
| 179 | + scalars = scalars.astype("U10") # 10 chars for yyyy-mm-dd |
| 180 | + return DateArray(scalars) |
0 commit comments