| 60 | |
| 61 | |
| 62 | class DateArray(ExtensionArray): |
| 63 | def __init__( |
| 64 | self, |
| 65 | dates: ( |
| 66 | dt.date |
| 67 | | Sequence[dt.date] |
| 68 | | tuple[np.ndarray, np.ndarray, np.ndarray] |
| 69 | | np.ndarray |
| 70 | ), |
| 71 | ) -> None: |
| 72 | if isinstance(dates, dt.date): |
| 73 | self._year = np.array([dates.year]) |
| 74 | self._month = np.array([dates.month]) |
| 75 | self._day = np.array([dates.year]) |
| 76 | return |
| 77 | |
| 78 | ldates = len(dates) |
| 79 | if isinstance(dates, list): |
| 80 | # pre-allocate the arrays since we know the size before hand |
| 81 | self._year = np.zeros(ldates, dtype=np.uint16) # 65535 (0, 9999) |
| 82 | self._month = np.zeros(ldates, dtype=np.uint8) # 255 (1, 31) |
| 83 | self._day = np.zeros(ldates, dtype=np.uint8) # 255 (1, 12) |
| 84 | # populate them |
| 85 | for i, (y, m, d) in enumerate( |
| 86 | (date.year, date.month, date.day) for date in dates |
| 87 | ): |
| 88 | self._year[i] = y |
| 89 | self._month[i] = m |
| 90 | self._day[i] = d |
| 91 | |
| 92 | elif isinstance(dates, tuple): |
| 93 | # only support triples |
| 94 | if ldates != 3: |
| 95 | raise ValueError("only triples are valid") |
| 96 | # check if all elements have the same type |
| 97 | if any(not isinstance(x, np.ndarray) for x in dates): |
| 98 | raise TypeError("invalid type") |
| 99 | ly, lm, ld = (len(cast(np.ndarray, d)) for d in dates) |
| 100 | if not ly == lm == ld: |
| 101 | raise ValueError( |
| 102 | f"tuple members must have the same length: {(ly, lm, ld)}" |
| 103 | ) |
| 104 | self._year = dates[0].astype(np.uint16) |
| 105 | self._month = dates[1].astype(np.uint8) |
| 106 | self._day = dates[2].astype(np.uint8) |
| 107 | |
| 108 | elif isinstance(dates, np.ndarray) and dates.dtype == "U10": |
| 109 | self._year = np.zeros(ldates, dtype=np.uint16) # 65535 (0, 9999) |
| 110 | self._month = np.zeros(ldates, dtype=np.uint8) # 255 (1, 31) |
| 111 | self._day = np.zeros(ldates, dtype=np.uint8) # 255 (1, 12) |
| 112 | |
| 113 | # error: "object_" object is not iterable |
| 114 | obj = np.char.split(dates, sep="-") |
| 115 | for (i,), (y, m, d) in np.ndenumerate(obj): |
| 116 | self._year[i] = int(y) |
| 117 | self._month[i] = int(m) |
| 118 | self._day[i] = int(d) |
| 119 |
no outgoing calls