increase coverage

py-polars/tests/unit/operations/namespaces/temporal/test_round.py

@@ -14,6 +14,8 @@
 if TYPE_CHECKING:
     from zoneinfo import ZoneInfo
 
+    from polars.type_aliases import TimeUnit
+
 
 else:
     from polars._utils.convert import string_to_zoneinfo as ZoneInfo
@@ -125,3 +127,65 @@ def test_dt_round_fast_path_vs_slow_path(datetimes: list[datetime], every: str)
     # Definitely uses slowpath:
     expected = s.dt.round(pl.Series([every] * len(datetimes)))
     assert_series_equal(result, expected)
+
+
+def test_round_date() -> None:
+    # n vs n
+    df = pl.DataFrame(
+        {"a": [date(2020, 1, 1), None, date(2020, 1, 19)], "b": [None, "1mo", "1mo"]}
+    )
+    result = df.select(pl.col("a").dt.round(pl.col("b")))["a"]
+    expected = pl.Series("a", [None, None, date(2020, 2, 1)])
+    assert_series_equal(result, expected)
+
+    # n vs 1
+    df = pl.DataFrame(
+        {"a": [date(2020, 1, 1), None, date(2020, 1, 3)], "b": [None, "1mo", "1mo"]}
+    )
+    result = df.select(pl.col("a").dt.round("1mo"))["a"]
+    expected = pl.Series("a", [date(2020, 1, 1), None, date(2020, 1, 1)])
+    assert_series_equal(result, expected)
+
+    # n vs missing
+    df = pl.DataFrame(
+        {"a": [date(2020, 1, 1), None, date(2020, 1, 3)], "b": [None, "1mo", "1mo"]}
+    )
+    result = df.select(pl.col("a").dt.round(pl.lit(None, dtype=pl.String)))["a"]
+    expected = pl.Series("a", [None, None, None], dtype=pl.Date)
+    assert_series_equal(result, expected)
+
+    # 1 vs n
+    df = pl.DataFrame(
+        {"a": [date(2020, 1, 1), None, date(2020, 1, 3)], "b": [None, "1mo", "1mo"]}
+    )
+    result = df.select(a=pl.date(2020, 1, 1).dt.round(pl.col("b")))["a"]
+    expected = pl.Series("a", [None, date(2020, 1, 1), date(2020, 1, 1)])
+    assert_series_equal(result, expected)
+
+    # missing vs n
+    df = pl.DataFrame(
+        {"a": [date(2020, 1, 1), None, date(2020, 1, 3)], "b": [None, "1mo", "1mo"]}
+    )
+    result = df.select(a=pl.lit(None, dtype=pl.Date).dt.round(pl.col("b")))["a"]
+    expected = pl.Series("a", [None, None, None], dtype=pl.Date)
+    assert_series_equal(result, expected)
+
+
+@pytest.mark.parametrize("time_unit", ["ms", "us", "ns"])
+def test_round_datetime_simple(time_unit: TimeUnit) -> None:
+    s = pl.Series([datetime(2020, 1, 2, 6)], dtype=pl.Datetime(time_unit))
+    result = s.dt.round("1mo").item()
+    assert result == datetime(2020, 1, 1)
+    result = s.dt.round("1d").item()
+    assert result == datetime(2020, 1, 2)
+
+
+@pytest.mark.parametrize("time_unit", ["ms", "us", "ns"])
+def test_round_datetime_w_expression(time_unit: TimeUnit) -> None:
+    df = pl.DataFrame(
+        {"a": [datetime(2020, 1, 2, 6), datetime(2020, 1, 20, 21)], "b": ["1mo", "1d"]},
+        schema_overrides={"a": pl.Datetime(time_unit)},
+    )
+    result = df.select(pl.col("a").dt.round(pl.col("b")))["a"]
+    assert result[0] == datetime(2020, 1, 1)
+    assert result[1] == datetime(2020, 1, 21)