From 8ba3a7fa0d99b73f7f43c1aed82c571b941eecea Mon Sep 17 00:00:00 2001 From: Mahmoud Hashemi Date: Wed, 7 Dec 2022 18:22:36 -0800 Subject: [PATCH] test and fix year step and december handling in date_range, fixes #319 --- boltons/timeutils.py | 10 ++++++---- tests/test_timeutils.py | 14 ++++++++++++++ 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/boltons/timeutils.py b/boltons/timeutils.py index 8d360d0..1ec31dd 100644 --- a/boltons/timeutils.py +++ b/boltons/timeutils.py @@ -377,6 +377,8 @@ def daterange(start, stop, step=1, inclusive=False): else: raise ValueError('step expected int, timedelta, or tuple' ' (year, month, day), not: %r' % step) + + m_step += y_step * 12 if stop is None: finished = lambda now, stop: False @@ -388,10 +390,10 @@ def daterange(start, stop, step=1, inclusive=False): while not finished(now, stop): yield now - if y_step or m_step: - m_y_step, cur_month = divmod(now.month + m_step, 12) - now = now.replace(year=now.year + y_step + m_y_step, - month=cur_month or 12) + if m_step: + m_y_step, cur_month = divmod((now.month - 1) + m_step, 12) + now = now.replace(year=now.year + m_y_step, + month=(cur_month + 1)) now = now + d_step return diff --git a/tests/test_timeutils.py b/tests/test_timeutils.py index 3e91ce6..631435e 100644 --- a/tests/test_timeutils.py +++ b/tests/test_timeutils.py @@ -48,6 +48,20 @@ def test_daterange_years(): assert years_from_2025[0] == date(2025, 1, 1) assert years_from_2025[-1] == date(2017, 1, 1) + +def test_daterange_years_step(): + start_day = date(year=2012, month=12, day=25) + end_day = date(year=2016, month=1, day=1) + dates = list(daterange(start_day, end_day, step=(1, 0, 0), inclusive=False)) + expected = [date(year=2012, month=12, day=25), date(year=2013, month=12, day=25), date(year=2014, month=12, day=25), date(year=2015, month=12, day=25)] + + assert dates == expected + + dates = list(daterange(start_day, end_day, step=(0, 13, 0), inclusive=False)) + expected = [date(year=2012, month=12, day=25), date(year=2014, month=1, day=25), date(year=2015, month=2, day=25)] + assert dates == expected + + def test_daterange_infinite(): today = date.today() infinite_dates = daterange(today, None)