1from __future__ import annotations
  2
  3import datetime
  4from calendar import monthrange
  5from collections.abc import Iterator
  6from enum import Enum
  7
  8from plain.utils import timezone
  9
 10__all__ = ["DatetimeRangeAliases", "DatetimeRange"]
 11
 12
 13class DatetimeRangeAliases(Enum):
 14    TODAY = "Today"
 15    THIS_WEEK = "This Week"
 16    THIS_WEEK_TO_DATE = "This Week-to-date"
 17    THIS_MONTH = "This Month"
 18    THIS_MONTH_TO_DATE = "This Month-to-date"
 19    THIS_QUARTER = "This Quarter"
 20    THIS_QUARTER_TO_DATE = "This Quarter-to-date"
 21    THIS_YEAR = "This Year"
 22    THIS_YEAR_TO_DATE = "This Year-to-date"
 23    LAST_WEEK = "Last Week"
 24    LAST_WEEK_TO_DATE = "Last Week-to-date"
 25    LAST_MONTH = "Last Month"
 26    LAST_MONTH_TO_DATE = "Last Month-to-date"
 27    LAST_QUARTER = "Last Quarter"
 28    LAST_QUARTER_TO_DATE = "Last Quarter-to-date"
 29    LAST_YEAR = "Last Year"
 30    LAST_YEAR_TO_DATE = "Last Year-to-date"
 31    SINCE_30_DAYS_AGO = "Since 30 Days Ago"
 32    SINCE_60_DAYS_AGO = "Since 60 Days Ago"
 33    SINCE_90_DAYS_AGO = "Since 90 Days Ago"
 34    SINCE_365_DAYS_AGO = "Since 365 Days Ago"
 35    NEXT_WEEK = "Next Week"
 36    NEXT_4_WEEKS = "Next 4 Weeks"
 37    NEXT_MONTH = "Next Month"
 38    NEXT_QUARTER = "Next Quarter"
 39    NEXT_YEAR = "Next Year"
 40
 41    # TODO doesn't include anything less than a day...
 42    # ex. SINCE_1_HOUR_AGO = "Since 1 Hour Ago"
 43
 44    def __str__(self) -> str:
 45        return self.value
 46
 47    @classmethod
 48    def from_value(cls, value: str) -> DatetimeRangeAliases:
 49        for member in cls:
 50            if member.value == value:
 51                return member
 52        raise ValueError(f"{value} is not a valid value for {cls.__name__}")
 53
 54    @classmethod
 55    def to_range(cls, value: str | DatetimeRangeAliases) -> DatetimeRange:
 56        # Convert enum to string value if needed
 57        if isinstance(value, DatetimeRangeAliases):
 58            value = value.value
 59
 60        now = timezone.localtime()
 61        start_of_today = now.replace(hour=0, minute=0, second=0, microsecond=0)
 62        start_of_week = start_of_today - datetime.timedelta(
 63            days=start_of_today.weekday()
 64        )
 65        start_of_month = start_of_today.replace(day=1)
 66        start_of_quarter = start_of_today.replace(
 67            month=((start_of_today.month - 1) // 3) * 3 + 1, day=1
 68        )
 69        start_of_year = start_of_today.replace(month=1, day=1)
 70
 71        def end_of_day(dt: datetime.datetime) -> datetime.datetime:
 72            return dt.replace(hour=23, minute=59, second=59, microsecond=999999)
 73
 74        def end_of_month(dt: datetime.datetime) -> datetime.datetime:
 75            last_day = monthrange(dt.year, dt.month)[1]
 76            return end_of_day(dt.replace(day=last_day))
 77
 78        def end_of_quarter(dt: datetime.datetime) -> datetime.datetime:
 79            end_month = ((dt.month - 1) // 3 + 1) * 3
 80            return end_of_month(dt.replace(month=end_month))
 81
 82        def end_of_year(dt: datetime.datetime) -> datetime.datetime:
 83            return end_of_month(dt.replace(month=12))
 84
 85        if value == cls.TODAY.value:
 86            return DatetimeRange(start_of_today, end_of_day(now))
 87        if value == cls.THIS_WEEK.value:
 88            return DatetimeRange(
 89                start_of_week, end_of_day(start_of_week + datetime.timedelta(days=6))
 90            )
 91        if value == cls.THIS_WEEK_TO_DATE.value:
 92            return DatetimeRange(start_of_week, now)
 93        if value == cls.THIS_MONTH.value:
 94            return DatetimeRange(start_of_month, end_of_month(start_of_month))
 95        if value == cls.THIS_MONTH_TO_DATE.value:
 96            return DatetimeRange(start_of_month, now)
 97        if value == cls.THIS_QUARTER.value:
 98            return DatetimeRange(start_of_quarter, end_of_quarter(start_of_quarter))
 99        if value == cls.THIS_QUARTER_TO_DATE.value:
100            return DatetimeRange(start_of_quarter, now)
101        if value == cls.THIS_YEAR.value:
102            return DatetimeRange(start_of_year, end_of_year(start_of_year))
103        if value == cls.THIS_YEAR_TO_DATE.value:
104            return DatetimeRange(start_of_year, now)
105        if value == cls.LAST_WEEK.value:
106            last_week_start = start_of_week - datetime.timedelta(days=7)
107            return DatetimeRange(
108                last_week_start,
109                end_of_day(last_week_start + datetime.timedelta(days=6)),
110            )
111        if value == cls.LAST_WEEK_TO_DATE.value:
112            return DatetimeRange(start_of_week - datetime.timedelta(days=7), now)
113        if value == cls.LAST_MONTH.value:
114            last_month = (start_of_month - datetime.timedelta(days=1)).replace(day=1)
115            return DatetimeRange(last_month, end_of_month(last_month))
116        if value == cls.LAST_MONTH_TO_DATE.value:
117            last_month = (start_of_month - datetime.timedelta(days=1)).replace(day=1)
118            return DatetimeRange(last_month, now)
119        if value == cls.LAST_QUARTER.value:
120            last_quarter = (start_of_quarter - datetime.timedelta(days=1)).replace(
121                day=1
122            )
123            return DatetimeRange(last_quarter, end_of_quarter(last_quarter))
124        if value == cls.LAST_QUARTER_TO_DATE.value:
125            last_quarter = (start_of_quarter - datetime.timedelta(days=1)).replace(
126                day=1
127            )
128            return DatetimeRange(last_quarter, now)
129        if value == cls.LAST_YEAR.value:
130            last_year = start_of_year.replace(year=start_of_year.year - 1)
131            return DatetimeRange(last_year, end_of_year(last_year))
132        if value == cls.LAST_YEAR_TO_DATE.value:
133            last_year = start_of_year.replace(year=start_of_year.year - 1)
134            return DatetimeRange(last_year, now)
135        if value == cls.SINCE_30_DAYS_AGO.value:
136            return DatetimeRange(now - datetime.timedelta(days=30), now)
137        if value == cls.SINCE_60_DAYS_AGO.value:
138            return DatetimeRange(now - datetime.timedelta(days=60), now)
139        if value == cls.SINCE_90_DAYS_AGO.value:
140            return DatetimeRange(now - datetime.timedelta(days=90), now)
141        if value == cls.SINCE_365_DAYS_AGO.value:
142            return DatetimeRange(now - datetime.timedelta(days=365), now)
143        if value == cls.NEXT_WEEK.value:
144            next_week_start = start_of_week + datetime.timedelta(days=7)
145            return DatetimeRange(
146                next_week_start,
147                end_of_day(next_week_start + datetime.timedelta(days=6)),
148            )
149        if value == cls.NEXT_4_WEEKS.value:
150            return DatetimeRange(now, end_of_day(now + datetime.timedelta(days=28)))
151        if value == cls.NEXT_MONTH.value:
152            next_month = (start_of_month + datetime.timedelta(days=31)).replace(day=1)
153            return DatetimeRange(next_month, end_of_month(next_month))
154        if value == cls.NEXT_QUARTER.value:
155            next_quarter = (start_of_quarter + datetime.timedelta(days=90)).replace(
156                day=1
157            )
158            return DatetimeRange(next_quarter, end_of_quarter(next_quarter))
159        if value == cls.NEXT_YEAR.value:
160            next_year = start_of_year.replace(year=start_of_year.year + 1)
161            return DatetimeRange(next_year, end_of_year(next_year))
162        raise ValueError(f"Invalid range: {value}")
163
164
165class DatetimeRange:
166    start: datetime.datetime
167    end: datetime.datetime
168
169    def __init__(
170        self,
171        start: datetime.datetime | datetime.date | str,
172        end: datetime.datetime | datetime.date | str,
173    ):
174        # Convert all inputs to datetime.datetime
175        if isinstance(start, str) and start:
176            self.start = datetime.datetime.fromisoformat(start)
177        elif isinstance(start, datetime.date) and not isinstance(
178            start, datetime.datetime
179        ):
180            self.start = timezone.localtime().replace(
181                year=start.year, month=start.month, day=start.day
182            )
183        else:
184            self.start = start  # type: ignore[assignment]
185
186        if isinstance(end, str) and end:
187            self.end = datetime.datetime.fromisoformat(end)
188        elif isinstance(end, datetime.date) and not isinstance(end, datetime.datetime):
189            self.end = timezone.localtime().replace(
190                year=end.year, month=end.month, day=end.day
191            )
192        else:
193            self.end = end  # type: ignore[assignment]
194
195    def as_tuple(self) -> tuple[datetime.datetime, datetime.datetime]:
196        return (self.start, self.end)
197
198    def total_days(self) -> int:
199        return (self.end - self.start).days
200
201    def iter_days(self) -> Iterator[datetime.date]:
202        """Yields each day in the range (inclusive of end date)."""
203        return iter(
204            self.start.date() + datetime.timedelta(days=i)
205            for i in range(0, self.total_days() + 1)
206        )
207
208    def iter_weeks(self) -> Iterator[datetime.datetime]:
209        """Yields the start of each week in the range."""
210        current = self.start - datetime.timedelta(days=self.start.weekday())
211        current = current.replace(hour=0, minute=0, second=0, microsecond=0)
212        while current <= self.end:
213            next_week = current + datetime.timedelta(weeks=1)
214            yield current
215            current = next_week
216
217    def iter_months(self) -> Iterator[datetime.datetime]:
218        """Yields the start of each month in the range."""
219        current = self.start.replace(day=1)
220        current = current.replace(hour=0, minute=0, second=0, microsecond=0)
221        while current <= self.end:
222            if current.month == 12:
223                next_month = current.replace(year=current.year + 1, month=1)
224            else:
225                next_month = current.replace(month=current.month + 1)
226            yield current
227            current = next_month
228
229    def iter_quarters(self) -> Iterator[datetime.datetime]:
230        """Yields the start of each quarter in the range."""
231        current = self.start.replace(month=((self.start.month - 1) // 3) * 3 + 1, day=1)
232        current = current.replace(hour=0, minute=0, second=0, microsecond=0)
233        while current <= self.end:
234            next_quarter_month = ((current.month - 1) // 3 + 1) * 3 + 1
235            if next_quarter_month > 12:
236                next_quarter_month -= 12
237                next_year = current.year + 1
238            else:
239                next_year = current.year
240            next_quarter = datetime.datetime(
241                next_year, next_quarter_month, 1, tzinfo=current.tzinfo
242            )
243            yield current
244            current = next_quarter
245
246    def iter_years(self) -> Iterator[datetime.datetime]:
247        """Yields the start of each year in the range."""
248        current = self.start.replace(month=1, day=1)
249        current = current.replace(hour=0, minute=0, second=0, microsecond=0)
250        while current <= self.end:
251            next_year = current.replace(year=current.year + 1)
252            yield current
253            current = next_year
254
255    def __repr__(self) -> str:
256        return f"DatetimeRange({self.start}, {self.end})"
257
258    def __str__(self) -> str:
259        return f"{self.start} to {self.end}"
260
261    def __eq__(self, other: object) -> bool:
262        if not isinstance(other, DatetimeRange):
263            return False
264        return self.start == other.start and self.end == other.end
265
266    def __hash__(self) -> int:
267        return hash((self.start, self.end))
268
269    def __contains__(self, item: datetime.datetime) -> bool:
270        return self.start <= item <= self.end