diff --git a/src/pendulum/interval.py b/src/pendulum/interval.py index c5e0713a7..2116ee967 100644 --- a/src/pendulum/interval.py +++ b/src/pendulum/interval.py @@ -300,6 +300,9 @@ def range(self, unit: str, amount: int = 1) -> Iterator[_T]: while op(start, end): yield start + if start == end and amount != 0: + break + start = getattr(self.start, method)(**{unit: i}) i += amount diff --git a/tests/interval/test_range.py b/tests/interval/test_range.py index 0b883e583..7d71b51c5 100644 --- a/tests/interval/test_range.py +++ b/tests/interval/test_range.py @@ -1,5 +1,9 @@ from __future__ import annotations +from itertools import islice + +import pytest + import pendulum from pendulum.interval import Interval @@ -117,3 +121,73 @@ def test_range_amount(): assert_datetime(r[1], 2016, 10, 16, 1, 0, 0) assert_datetime(r[2], 2016, 10, 18, 0, 0, 0) assert_datetime(r[3], 2016, 10, 20, 0, 0, 0) + + +@pytest.mark.parametrize("factory", [pendulum.date, pendulum.datetime]) +@pytest.mark.parametrize("absolute", [False, True]) +@pytest.mark.parametrize( + "start, end, unit, amount", + [ + ((9999, 12, 30), (9999, 12, 31), "days", 1), + ((1, 1, 2), (1, 1, 1), "days", 1), + ((9999, 12, 31), (9999, 12, 31), "days", 1), + ((1, 1, 1), (1, 1, 1), "days", 1), + ((9999, 12, 29), (9999, 12, 31), "days", 2), + ((1, 1, 3), (1, 1, 1), "days", 2), + ((9999, 12, 24), (9999, 12, 31), "weeks", 1), + ((1, 1, 8), (1, 1, 1), "weeks", 1), + ((9999, 11, 30), (9999, 12, 30), "months", 1), + ((1, 2, 1), (1, 1, 1), "months", 1), + ((9998, 12, 31), (9999, 12, 31), "years", 1), + ((2, 1, 1), (1, 1, 1), "years", 1), + ], +) +def test_range_stops_at_final_endpoint(factory, absolute, start, end, unit, amount): + start, end = factory(*start), factory(*end) + interval = pendulum.interval(start, end, absolute=absolute) + expected = [start] if start == end else [start, end] + if absolute: + expected.sort() + + assert list(interval.range(unit, amount)) == expected + + +@pytest.mark.parametrize("reverse", [False, True]) +@pytest.mark.parametrize("unit", ["hours", "minutes", "seconds", "microseconds"]) +def test_range_stops_at_final_datetime_endpoint(reverse, unit): + if reverse: + end = pendulum.datetime(1, 1, 1) + start = end.add(**{unit: 1}) + else: + end = pendulum.datetime(9999, 12, 31, 23, 59, 59, 999999) + start = end.subtract(**{unit: 1}) + + assert list(pendulum.interval(start, end).range(unit)) == [start, end] + + +@pytest.mark.parametrize("end_day", [1, 3]) +def test_range_zero_amount(end_day): + start = pendulum.date(2000, 1, 1) + end = pendulum.date(2000, 1, end_day) + + assert list(islice(pendulum.interval(start, end).range("days", 0), 3)) == [ + start, + start, + start, + ] + + +@pytest.mark.parametrize("reverse", [False, True]) +@pytest.mark.parametrize("absolute", [False, True]) +def test_range_amount_does_not_reach_endpoint(reverse, absolute): + start, end = pendulum.date(2000, 1, 1), pendulum.date(2000, 1, 4) + if reverse: + start, end = end, start + expected = [pendulum.date(2000, 1, 1), pendulum.date(2000, 1, 3)] + if reverse and not absolute: + expected = [pendulum.date(2000, 1, 4), pendulum.date(2000, 1, 2)] + + assert ( + list(pendulum.interval(start, end, absolute=absolute).range("days", 2)) + == expected + )