fix: Fix behavior of month intervals in date_range (#12317)

crates/polars-time/src/date_range.rs

@@ -132,32 +132,25 @@ pub(crate) fn datetime_range_i64(
     }
     let mut ts = Vec::with_capacity(size);
 
-    let mut t = start;
+    let mut i = match closed {
+        ClosedWindow::Both | ClosedWindow::Left => 0,
+        ClosedWindow::Right | ClosedWindow::None => 1,
+    };
+    let mut t = offset_fn(&(interval * i), start, tz)?;
+    i += 1;
     match closed {
-        ClosedWindow::Both => {
-            while t <= end {
-                ts.push(t);
-                t = offset_fn(&interval, t, tz)?
-            }
-        },
-        ClosedWindow::Left => {
-            while t < end {
-                ts.push(t);
-                t = offset_fn(&interval, t, tz)?
-            }
-        },
-        ClosedWindow::Right => {
-            t = offset_fn(&interval, t, tz)?;
+        ClosedWindow::Both | ClosedWindow::Right => {
             while t <= end {
                 ts.push(t);
-                t = offset_fn(&interval, t, tz)?
+                t = offset_fn(&(interval * i), start, tz)?;
+                i += 1;
             }
         },
-        ClosedWindow::None => {
-            t = offset_fn(&interval, t, tz)?;
+        ClosedWindow::Left | ClosedWindow::None => {
             while t < end {
                 ts.push(t);
-                t = offset_fn(&interval, t, tz)?
+                t = offset_fn(&(interval * i), start, tz)?;
+                i += 1;
             }
         },
     }

py-polars/polars/functions/range/date_range.py

@@ -95,9 +95,6 @@ def date_range(
     interval
         Interval of the range periods, specified as a Python `timedelta` object
         or using the Polars duration string language (see "Notes" section below).
-
-        To create a month-end date series, combine with :meth:`Expr.dt.month_end` (see
-        "Examples" section below).
     closed : {'both', 'left', 'right', 'none'}
         Define which sides of the range are closed (inclusive).
     time_unit : {None, 'ns', 'us', 'ms'}
@@ -184,19 +181,6 @@ def date_range(
         1985-01-09
     ]
 
-    Combine with :meth:`Expr.dt.month_end` to get the last day of the month:
-
-    >>> pl.date_range(
-    ...     date(2022, 1, 1), date(2022, 3, 1), "1mo", eager=True
-    ... ).dt.month_end()
-    shape: (3,)
-    Series: 'date' [date]
-    [
-        2022-01-31
-        2022-02-28
-        2022-03-31
-    ]
-
     """
     interval = deprecate_saturating(interval)
     if name is not None:

py-polars/tests/unit/functions/range/test_date_range.py

@@ -10,7 +10,7 @@
 from polars.testing import assert_frame_equal, assert_series_equal
 
 if TYPE_CHECKING:
-    from polars.type_aliases import TimeUnit
+    from polars.type_aliases import ClosedInterval, TimeUnit
 
 
 def test_date_range() -> None:
@@ -138,28 +138,22 @@ def test_date_ranges_single_row_lazy_7110() -> None:
     assert_frame_equal(result, expected)
 
 
-def test_date_range_end_of_month_5441() -> None:
+@pytest.mark.parametrize(
+    ("closed", "expected_values"),
+    [
+        ("right", [date(2020, 2, 29), date(2020, 3, 31)]),
+        ("left", [date(2020, 1, 31), date(2020, 2, 29)]),
+        ("none", [date(2020, 2, 29)]),
+        ("both", [date(2020, 1, 31), date(2020, 2, 29), date(2020, 3, 31)]),
+    ],
+)
+def test_date_range_end_of_month_5441(
+    closed: ClosedInterval, expected_values: list[date]
+) -> None:
     start = date(2020, 1, 31)
-    stop = date(2021, 1, 31)
-    result = pl.date_range(start, stop, interval="1mo", eager=True)
-    expected = pl.Series(
-        "date",
-        [
-            date(2020, 1, 31),
-            date(2020, 2, 29),
-            date(2020, 3, 29),
-            date(2020, 4, 29),
-            date(2020, 5, 29),
-            date(2020, 6, 29),
-            date(2020, 7, 29),
-            date(2020, 8, 29),
-            date(2020, 9, 29),
-            date(2020, 10, 29),
-            date(2020, 11, 29),
-            date(2020, 12, 29),
-            date(2021, 1, 29),
-        ],
-    )
+    stop = date(2020, 3, 31)
+    result = pl.date_range(start, stop, interval="1mo", closed=closed, eager=True)
+    expected = pl.Series("date", expected_values)
     assert_series_equal(result, expected)
 
 

py-polars/tests/unit/functions/range/test_datetime_range.py

@@ -8,13 +8,13 @@
 import polars as pl
 from polars.datatypes import DTYPE_TEMPORAL_UNITS
 from polars.exceptions import ComputeError, TimeZoneAwareConstructorWarning
-from polars.testing import assert_frame_equal
+from polars.testing import assert_frame_equal, assert_series_equal
 
 if TYPE_CHECKING:
     from zoneinfo import ZoneInfo
 
     from polars.datatypes import PolarsDataType
-    from polars.type_aliases import TimeUnit
+    from polars.type_aliases import ClosedInterval, TimeUnit
 else:
     from polars.utils.convert import get_zoneinfo as ZoneInfo
 
@@ -475,3 +475,22 @@ def test_datetime_range_invalid_interval(interval: timedelta) -> None:
         pl.datetime_range(
             datetime(2000, 3, 20), datetime(2000, 3, 21), interval="-1h", eager=True
         )
+
+
+@pytest.mark.parametrize(
+    ("closed", "expected_values"),
+    [
+        ("right", [datetime(2020, 2, 29), datetime(2020, 3, 31)]),
+        ("left", [datetime(2020, 1, 31), datetime(2020, 2, 29)]),
+        ("none", [datetime(2020, 2, 29)]),
+        ("both", [datetime(2020, 1, 31), datetime(2020, 2, 29), datetime(2020, 3, 31)]),
+    ],
+)
+def test_datetime_range_end_of_month_5441(
+    closed: ClosedInterval, expected_values: list[datetime]
+) -> None:
+    start = date(2020, 1, 31)
+    stop = date(2020, 3, 31)
+    result = pl.datetime_range(start, stop, interval="1mo", closed=closed, eager=True)
+    expected = pl.Series("datetime", expected_values)
+    assert_series_equal(result, expected)