From 53cc46d5f50035db809c72d807312eb68538c7fa Mon Sep 17 00:00:00 2001 From: el-g-1 Date: Mon, 13 Sep 2021 22:27:27 -0700 Subject: [PATCH] TST: Add test for masking timedelta (#39548) --- pandas/tests/frame/indexing/test_mask.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/pandas/tests/frame/indexing/test_mask.py b/pandas/tests/frame/indexing/test_mask.py index ac80426883dd5..70ec02a2334af 100644 --- a/pandas/tests/frame/indexing/test_mask.py +++ b/pandas/tests/frame/indexing/test_mask.py @@ -9,6 +9,7 @@ DataFrame, Series, StringDtype, + Timedelta, isna, ) import pandas._testing as tm @@ -136,3 +137,16 @@ def test_mask_stringdtype(): dtype=StringDtype(), ) tm.assert_frame_equal(result, expected) + + +def test_mask_where_dtype_timedelta(): + # https://github.com/pandas-dev/pandas/issues/39548 + df = DataFrame([Timedelta(i, unit="d") for i in range(5)]) + + expected = DataFrame(np.full(5, np.nan, dtype="timedelta64[ns]")) + tm.assert_frame_equal(df.mask(df.notna()), expected) + + expected = DataFrame( + [np.nan, np.nan, np.nan, Timedelta("3 day"), Timedelta("4 day")] + ) + tm.assert_frame_equal(df.where(df > Timedelta(2, unit="d")), expected)