|
1 | 1 | from enum import IntEnum |
| 2 | +from dataclasses import dataclass, replace |
| 3 | +import operator |
2 | 4 |
|
3 | 5 | # Pre 3.10 requires Union for multiple types, e.g. Union[int, None] instead of int | None |
4 | | -from typing import Optional, Union |
| 6 | +from typing import Optional, Union, Iterable |
5 | 7 |
|
6 | 8 | import numpy as np |
7 | 9 |
|
@@ -29,6 +31,139 @@ def days(self) -> int: |
29 | 31 | return int(self.astype("datetime64[D]").astype("int")) |
30 | 32 |
|
31 | 33 |
|
| 34 | +@dataclass |
| 35 | +class UnInt: |
| 36 | + lower: int |
| 37 | + upper: int |
| 38 | + |
| 39 | + def __post_init__(self): |
| 40 | + # validate that lower value is less than upper |
| 41 | + if not self.lower < self.upper: |
| 42 | + raise ValueError( |
| 43 | + f"Lower value ({self.lower}) must be less than upper ({self.upper})" |
| 44 | + ) |
| 45 | + |
| 46 | + def __iter__(self) -> Iterable: |
| 47 | + # yield all integers in range from lower to upper, inclusive |
| 48 | + yield from range(self.lower, self.upper + 1) |
| 49 | + |
| 50 | + def __gt__(self, other: object) -> bool: |
| 51 | + match other: |
| 52 | + case int(): |
| 53 | + return self.upper > other |
| 54 | + case UnInt(): |
| 55 | + return self.upper > other.lower |
| 56 | + case _: |
| 57 | + return NotImplemented |
| 58 | + |
| 59 | + def __lt__(self, other: object) -> bool: |
| 60 | + match other: |
| 61 | + case int(): |
| 62 | + return self.upper < other |
| 63 | + case UnInt(): |
| 64 | + return self.upper < other.lower |
| 65 | + case _: |
| 66 | + return NotImplemented |
| 67 | + |
| 68 | + def __contains__(self, other: object) -> bool: |
| 69 | + match other: |
| 70 | + case int(): |
| 71 | + return other >= self.lower and other <= self.upper |
| 72 | + case UnInt(): |
| 73 | + return other.lower >= self.lower and other.upper <= self.upper |
| 74 | + case _: |
| 75 | + # unsupported type: return false |
| 76 | + return False |
| 77 | + |
| 78 | + def _replace_with(self, other_lower, other_upper, op): |
| 79 | + """Create and return a new instance of UnInt using the specified |
| 80 | + operator (e.g. add, subtract) and other values to modify the values in |
| 81 | + the current UnInt instance.""" |
| 82 | + return replace( |
| 83 | + self, lower=op(self.lower, other_lower), upper=op(self.upper, other_upper) |
| 84 | + ) |
| 85 | + |
| 86 | + def __add__(self, other: object) -> bool: |
| 87 | + match other: |
| 88 | + case int(): |
| 89 | + # increase both values by the added amount |
| 90 | + add_values = (other, other) |
| 91 | + case UnInt(): |
| 92 | + # subtract the upper and lower values by the other lower and upper |
| 93 | + # to include the largest range of possible values |
| 94 | + # (when calculating with uncertain values, the uncertainty increases) |
| 95 | + add_values = (other.lower, other.upper) |
| 96 | + case _: |
| 97 | + return NotImplemented |
| 98 | + |
| 99 | + return self._replace_with(*add_values, operator.add) |
| 100 | + |
| 101 | + def __sub__(self, other): |
| 102 | + match other: |
| 103 | + case int(): |
| 104 | + # decrease both values by the subtracted amount |
| 105 | + sub_values = (other, other) |
| 106 | + case UnInt(): |
| 107 | + # to determine the largest range of possible values, |
| 108 | + # subtract the other upper value from current lower |
| 109 | + # and other lower value from current upper |
| 110 | + sub_values = (other.upper, other.lower) |
| 111 | + case _: |
| 112 | + return NotImplemented |
| 113 | + |
| 114 | + return self._replace_with(*sub_values, operator.sub) |
| 115 | + |
| 116 | + |
| 117 | +@dataclass |
| 118 | +class UnDelta: |
| 119 | + """ |
| 120 | + An uncertain timedelta, for durations where the number of days is uncertain. |
| 121 | + Initialize with a list of possible durations in days as integers, which are used |
| 122 | + to calculate a value for duration in :attr:`days` as an |
| 123 | + instance of :class:`UnInt`. |
| 124 | + """ |
| 125 | + |
| 126 | + # NOTE: we will probably need other timedelta-like logic here besides days... |
| 127 | + |
| 128 | + #: possible durations days, as an instance of :class:`UnInt` |
| 129 | + days: UnInt |
| 130 | + |
| 131 | + def __init__(self, *days: int): |
| 132 | + if len(days) < 2: |
| 133 | + raise ValueError( |
| 134 | + "Must specify at least two values for an uncertain duration" |
| 135 | + ) |
| 136 | + self.days = UnInt(min(days), max(days)) |
| 137 | + |
| 138 | + def __repr__(self): |
| 139 | + # customize string representation for simpler notation; default |
| 140 | + # specifies full UnInt initialization with upper and lower keywords |
| 141 | + return f"{self.__class__.__name__}(days=[{self.days.lower},{self.days.upper}])" |
| 142 | + |
| 143 | + # TODO: what does equality for an uncertain range mean? |
| 144 | + # is an uncertain range ever equal to another uncertain range? |
| 145 | + |
| 146 | + def __eq__(self, other: object) -> bool: |
| 147 | + # is an uncertain duration ever *equal* another, even if the values are the same? |
| 148 | + if other is self: |
| 149 | + return True |
| 150 | + return False |
| 151 | + |
| 152 | + def __lt__(self, other: object) -> bool: |
| 153 | + match other: |
| 154 | + case Timedelta() | UnDelta(): |
| 155 | + return self.days < other.days |
| 156 | + case _: |
| 157 | + return NotImplemented |
| 158 | + |
| 159 | + def __gt__(self, other: object) -> bool: |
| 160 | + match other: |
| 161 | + case Timedelta() | UnDelta(): |
| 162 | + return self.days > other.days |
| 163 | + case _: |
| 164 | + return NotImplemented |
| 165 | + |
| 166 | + |
32 | 167 | #: timedelta for single day |
33 | 168 | ONE_DAY = Timedelta(1) # ~ equivalent to datetime.timedelta(days=1) |
34 | 169 | #: timedelta for a single year (non-leap year) |
|
0 commit comments