diff --git a/arrow/arrow.py b/arrow/arrow.py index eecf23266..ae981fb59 100644 --- a/arrow/arrow.py +++ b/arrow/arrow.py @@ -644,7 +644,7 @@ def span_range( cls, frame: _T_FRAMES, start: dt_datetime, - end: dt_datetime, + end: Optional[dt_datetime] = None, tz: Optional[TZ_EXPR] = None, limit: Optional[int] = None, bounds: _BOUNDS = "[)", @@ -703,7 +703,7 @@ def span_range( tzinfo = cls._get_tzinfo(start.tzinfo if tz is None else tz) start = cls.fromdatetime(start, tzinfo).span(frame, exact=exact)[0] - end = cls.fromdatetime(end, tzinfo) + range_end = cls.fromdatetime(end, tzinfo) if end is not None else None _range = cls.range(frame, start, end, tz, limit) if not exact: for r in _range: @@ -711,13 +711,13 @@ def span_range( for r in _range: floor, ceil = r.span(frame, bounds=bounds, exact=exact) - if ceil > end: - ceil = end + if range_end is not None and ceil > range_end: + ceil = range_end if bounds[1] == ")": ceil += relativedelta(microseconds=-1) - if floor == end: + if floor == range_end: break - elif floor + relativedelta(microseconds=-1) == end: + elif floor + relativedelta(microseconds=-1) == range_end: break yield floor, ceil diff --git a/tests/test_arrow.py b/tests/test_arrow.py index b595e4e21..4bef15027 100644 --- a/tests/test_arrow.py +++ b/tests/test_arrow.py @@ -1448,6 +1448,68 @@ def test_exact_bound_exclude(self): assert result == expected + def test_limit_without_end(self): + # ``end`` is documented as optional: passing ``limit`` alone must + # return at most ``limit`` timespans, counted from ``start``. + expected = [ + (arrow.Arrow(2013, 1, 1), arrow.Arrow(2013, 1, 31, 23, 59, 59, 999999)), + (arrow.Arrow(2013, 2, 1), arrow.Arrow(2013, 2, 28, 23, 59, 59, 999999)), + (arrow.Arrow(2013, 3, 1), arrow.Arrow(2013, 3, 31, 23, 59, 59, 999999)), + ] + + assert ( + list(arrow.Arrow.span_range("month", datetime(2013, 1, 1), limit=3)) + == expected + ) + assert ( + list(arrow.Arrow.span_range("month", datetime(2013, 1, 1), None, limit=3)) + == expected + ) + + def test_limit_without_end_quarter_and_tz(self): + result = list( + arrow.Arrow.span_range( + "quarter", datetime(2013, 2, 2), tz="US/Pacific", limit=2 + ) + ) + + assert result == [ + ( + arrow.Arrow(2013, 1, 1, tzinfo=ZoneInfo("US/Pacific")), + arrow.Arrow( + 2013, 3, 31, 23, 59, 59, 999999, tzinfo=ZoneInfo("US/Pacific") + ), + ), + ( + arrow.Arrow(2013, 4, 1, tzinfo=ZoneInfo("US/Pacific")), + arrow.Arrow( + 2013, 6, 30, 23, 59, 59, 999999, tzinfo=ZoneInfo("US/Pacific") + ), + ), + ] + + def test_limit_without_end_exact(self): + result = list( + arrow.Arrow.span_range( + "month", datetime(2013, 1, 15, 12, 30), limit=2, exact=True + ) + ) + + assert result == [ + ( + arrow.Arrow(2013, 1, 15, 12, 30), + arrow.Arrow(2013, 2, 15, 12, 29, 59, 999999), + ), + ( + arrow.Arrow(2013, 2, 15, 12, 30), + arrow.Arrow(2013, 3, 15, 12, 29, 59, 999999), + ), + ] + + def test_limit_or_end_is_required(self): + with pytest.raises(ValueError): + list(arrow.Arrow.span_range("month", datetime(2013, 1, 1))) + def test_exact_floor_equals_end(self): result = list( arrow.Arrow.span_range(