matemat/matemat/util/monthdelta.py

32 lines
1.3 KiB
Python

from typing import Tuple
from datetime import datetime, timedelta
import calendar
def add_months(d: datetime, months: int) -> datetime:
"""
Add the given number of months to the passed date, considering the varying numbers of days in a month.
:param d: The date time to add to.
:param months: The number of months to add to.
:return: A datetime object offset by the requested number of months.
"""
if not isinstance(d, datetime) or not isinstance(months, int):
raise TypeError()
if months < 0:
raise ValueError('Can only add a positive number of months.')
nextmonth: Tuple[int, int] = (d.year, d.month)
days: int = 0
# Iterate the months between the passed date and the target month
for i in range(months):
days += calendar.monthrange(*nextmonth)[1]
if nextmonth[1] == 12:
nextmonth = nextmonth[0] + 1, 1
else:
nextmonth = nextmonth[0], nextmonth[1] + 1
# Set the day of month temporarily to 1, then add the day offset to reach the 1st of the target month
newdate: datetime = d.replace(day=1) + timedelta(days=days)
# Re-set the day of month to the intended value, but capped by the max. day in the target month
newdate = newdate.replace(day=min(d.day, calendar.monthrange(newdate.year, newdate.month)[1]))
return newdate