fix date_range time-unit kwargs (#2298)

pola-rs · Jan 7, 2022 · 78aa370 · 78aa370
1 parent f38e3e8
commit 78aa370
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/py-polars/polars/internals/functions.py b/py-polars/polars/internals/functions.py
@@ -230,6 +230,8 @@ def date_range(
 
     if in_nanoseconds_window(low) and in_nanoseconds_window(high) and time_unit is None:
         tu = "ns"
+    elif time_unit is not None:
+        tu = time_unit
     else:
         tu = "ms"
 

diff --git a/py-polars/tests/test_datelike.py b/py-polars/tests/test_datelike.py
@@ -236,6 +236,15 @@ def test_date_range() -> None:
     assert result.dt[2] == datetime(1985, 1, 4, 0, 0)
     assert result.dt[-1] == datetime(2015, 6, 30, 12, 0)
 
+    for tu in ["ns", "ms"]:
+        rng = pl.date_range(
+            datetime(2020, 1, 1), datetime(2020, 1, 2), "2h", time_unit=tu
+        )
+        assert rng.time_unit == tu
+        assert rng.shape == (13,)
+        assert rng.dt[0] == datetime(2020, 1, 1)
+        assert rng.dt[-1] == datetime(2020, 1, 2)
+
 
 def test_date_comp() -> None:
     one = datetime(2001, 1, 1)