GH648 Allowing dateoffset weekday from relativedelta (pandas-dev#1010)

* GH648 Allowing dateoffset weekday from relativedelta * Fix lint * GH648 Clean up testing * PR Feedback
loicdiridollou · Oct 2, 2024 · 237977f · 237977f
1 parent 1bc27e6
commit 237977f
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 1 deletion.
diff --git a/pandas-stubs/_libs/tslibs/offsets.pyi b/pandas-stubs/_libs/tslibs/offsets.pyi
@@ -12,6 +12,7 @@ from typing import (
     overload,
 )
 
+from dateutil.relativedelta import weekday as WeekdayClass
 import numpy as np
 from pandas.core.indexes.datetimes import DatetimeIndex
 from typing_extensions import Self
@@ -257,7 +258,7 @@ class DateOffset(RelativeDeltaOffset):
         year: int = ...,
         month: int = ...,
         day: int = ...,
-        weekday: int = ...,
+        weekday: int | WeekdayClass = ...,
         hour: int = ...,
         minute: int = ...,
         second: int = ...,

diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py
@@ -8,6 +8,15 @@
     cast,
 )
 
+from dateutil.relativedelta import (
+    FR,
+    MO,
+    SA,
+    SU,
+    TH,
+    TU,
+    WE,
+)
 import numpy as np
 from numpy import typing as npt
 import pandas as pd
@@ -1284,6 +1293,42 @@ def test_weekofmonth_init():
     )
 
 
+def test_dateoffset_weekday() -> None:
+    """Check that you can create a `pd.DateOffset` from weekday of int or relativedelta.weekday."""
+    check(
+        assert_type(pd.offsets.DateOffset(weekday=1), pd.offsets.DateOffset),
+        pd.offsets.DateOffset,
+    )
+    check(
+        assert_type(pd.offsets.DateOffset(weekday=MO), pd.offsets.DateOffset),
+        pd.offsets.DateOffset,
+    )
+    check(
+        assert_type(pd.offsets.DateOffset(weekday=TU), pd.offsets.DateOffset),
+        pd.offsets.DateOffset,
+    )
+    check(
+        assert_type(pd.offsets.DateOffset(weekday=WE), pd.offsets.DateOffset),
+        pd.offsets.DateOffset,
+    )
+    check(
+        assert_type(pd.offsets.DateOffset(weekday=TH), pd.offsets.DateOffset),
+        pd.offsets.DateOffset,
+    )
+    check(
+        assert_type(pd.offsets.DateOffset(weekday=FR), pd.offsets.DateOffset),
+        pd.offsets.DateOffset,
+    )
+    check(
+        assert_type(pd.offsets.DateOffset(weekday=SA), pd.offsets.DateOffset),
+        pd.offsets.DateOffset,
+    )
+    check(
+        assert_type(pd.offsets.DateOffset(weekday=SU), pd.offsets.DateOffset),
+        pd.offsets.DateOffset,
+    )
+
+
 def test_date_range_unit():
     check(
         assert_type(