import calendar
import datetime
from functools import lru_cache
from dataclasses import dataclass
import pytz
AVG_MONTH_LEN = (365 * 3 + 366) / (4 * 12)
[docs]
def midnight(dt: datetime.datetime):
"""
Get datetime of ``dt`` at midnight.
"""
return dt.replace(hour=0, minute=0, second=0, microsecond=0)
[docs]
@dataclass()
class MonthOfDateRange:
index: int
start: datetime.datetime
end: datetime.datetime
[docs]
def today_in_timezone(tz: datetime.timezone | None) -> datetime.datetime:
"""
Get datetime representing start of current day in the given timezone.
:param tz: Timezone. If not given, UTC will be used.
:return: Start of the current day.
"""
if tz is None:
tz = datetime.UTC
now = datetime.datetime.now(tz)
today = now.replace(hour=0, minute=0, second=0, microsecond=0)
return today
[docs]
def add_months(
dt: datetime.date | datetime.datetime, months: int, keep_last_day=False
) -> datetime.date | datetime.datetime:
"""
Add the given number of months to the given date or datetime. If the resulting month has less days than the input
month, last day of the month will be used.
If you set ``keep_last_day=True``, the resulting month has more days, and the input date was the last
day of the month, the last day of the resulting month will be used.
Examples:
>>> from datetime import date
>>> add_months(date(2020, 1, 31), months=1)
date(2020, 2, 28)
>>> add_months(date(2020, 1, 31), months=2)
date(2020, 3, 31)
>>> add_months(date(2020, 2, 28), months=1, keep_last_day=False)
date(2020, 3, 28)
>>> add_months(date(2020, 2, 28), months=1, keep_last_day=True)
date(2020, 3, 31)
:param dt: Input date or datetime.
:param months: Number of months to add.
:param keep_last_day: If the input date was the last day of the month, use the last day of month in the output.
:return: Resulting date or datetime.
"""
year = dt.year
month = dt.month
day = dt.day
_, current_ndays = calendar.monthrange(year, month)
month += months
if months > 0:
while month > 12:
year += 1
month -= 12
else:
while month < 1:
year -= 1
month += 12
_, new_ndays = calendar.monthrange(year, month)
if day > new_ndays:
day = new_ndays
elif keep_last_day and day == current_ndays:
day = new_ndays
return dt.replace(year=year, month=month, day=day)
[docs]
def get_previous_period(start_date: datetime.date, end_date: datetime.date) -> tuple[datetime.date, datetime.date]:
"""
Takes date range and creates a new date range with the exact number of days which precedes it. E.g.
For 2019-01-10 - 2019-01-18, this function will return 2019-01-01 - 2019-01-09
:param start_date: Start date.
:param end_date: End date.
:return: Date range that precedes the given one.
"""
total_delta = end_date - start_date
return (
start_date - datetime.timedelta(days=total_delta.days + 1),
start_date - datetime.timedelta(days=1),
)
[docs]
def get_month_in_date_range(
dt: datetime.datetime, date_range_start: datetime.datetime, date_range_end: datetime.datetime, keep_last_day=True
) -> MonthOfDateRange:
"""
Helper function for determining current month in a date range.
:param dt: Point in time inside the date range.
:param date_range_start: Start of the date range.
:param date_range_end: End of the date range.
:param keep_last_day: Keep the last day of month when determining end of months.
:return: Info about the current month.
"""
if dt < date_range_start or dt > date_range_end:
raise ValueError(
f'Date {dt!r} does not belong to the specified date range <{date_range_start!r}; {date_range_end!r}>.'
)
month_start = date_range_start
index = 0
while month_start < date_range_end:
month_end = add_months(month_start, months=1, keep_last_day=keep_last_day)
is_last = month_end >= date_range_end
if dt >= month_start and (dt < month_end or (is_last and dt <= month_end)):
return MonthOfDateRange(
index=index,
start=month_start,
end=min(date_range_end, month_end),
)
index += 1
month_start = month_end
raise ValueError('Could not determine the month in date range!')
[docs]
@lru_cache(maxsize=64)
def tzinfo(time_zone: str) -> datetime.tzinfo:
"""
Cached :class:`pytz.timezone`.
:param time_zone: timezone identifier
:return: datetime.tzinfo
:raises pytz.UnknownTimeZoneError in case of ``time_zone`` is unknown timezone identifier.
"""
return pytz.timezone(time_zone)