37 lines
1.3 KiB
Python
37 lines
1.3 KiB
Python
from datetime import datetime, timedelta, timezone
|
|
from typing import Optional
|
|
|
|
|
|
def add_months(base_dt: datetime, months_to_add: int) -> datetime:
|
|
"""Add calendar months to a datetime, clamping the day to the month's length.
|
|
|
|
Preserves tzinfo from base_dt.
|
|
"""
|
|
year = base_dt.year
|
|
month = base_dt.month + months_to_add
|
|
day = base_dt.day
|
|
|
|
# Normalize year and month
|
|
year += (month - 1) // 12
|
|
month = ((month - 1) % 12) + 1
|
|
|
|
# Determine last day of target month by rolling to next month's first day and subtracting 1 day
|
|
if month == 12:
|
|
next_month_first = datetime(year + 1, 1, 1, tzinfo=base_dt.tzinfo)
|
|
else:
|
|
next_month_first = datetime(year, month + 1, 1, tzinfo=base_dt.tzinfo)
|
|
last_day = (next_month_first - timedelta(days=1)).day
|
|
|
|
clamped_day = min(day, last_day)
|
|
return base_dt.replace(year=year, month=month, day=clamped_day)
|
|
|
|
|
|
def month_start(base_dt: Optional[datetime] = None) -> datetime:
|
|
"""Return the first instant of the month in UTC for a datetime."""
|
|
moment = base_dt or datetime.now(timezone.utc)
|
|
if moment.tzinfo is None:
|
|
moment = moment.replace(tzinfo=timezone.utc)
|
|
else:
|
|
moment = moment.astimezone(timezone.utc)
|
|
return datetime(moment.year, moment.month, 1, tzinfo=timezone.utc)
|