Skip to content

Commit a248f72

Browse files
authored
REF: Make WeekOfMonthMixin a cdef class (#34265)
1 parent af66d64 commit a248f72

File tree

4 files changed

+95
-3
lines changed

4 files changed

+95
-3
lines changed

doc/source/reference/offset_frequency.rst

+2
Original file line numberDiff line numberDiff line change
@@ -549,6 +549,7 @@ Methods
549549
WeekOfMonth.is_anchored
550550
WeekOfMonth.is_on_offset
551551
WeekOfMonth.__call__
552+
WeekOfMonth.weekday
552553

553554
LastWeekOfMonth
554555
---------------
@@ -569,6 +570,7 @@ Properties
569570
LastWeekOfMonth.normalize
570571
LastWeekOfMonth.rule_code
571572
LastWeekOfMonth.n
573+
LastWeekOfMonth.weekday
572574

573575
Methods
574576
~~~~~~~

pandas/_libs/tslibs/offsets.pyx

+7-2
Original file line numberDiff line numberDiff line change
@@ -562,6 +562,7 @@ cdef class BaseOffset:
562562
exclude = {"n", "inc", "normalize"}
563563
attrs = []
564564
for attr in sorted(self._attributes):
565+
# _attributes instead of __dict__ because cython attrs are not in __dict__
565566
if attr.startswith("_") or attr == "kwds" or not hasattr(self, attr):
566567
# DateOffset may not have some of these attributes
567568
continue
@@ -1164,13 +1165,17 @@ class CustomMixin:
11641165
object.__setattr__(self, "calendar", calendar)
11651166

11661167

1167-
class WeekOfMonthMixin(SingleConstructorOffset):
1168+
cdef class WeekOfMonthMixin(SingleConstructorOffset):
11681169
"""
11691170
Mixin for methods common to WeekOfMonth and LastWeekOfMonth.
11701171
"""
1172+
1173+
cdef readonly:
1174+
int weekday
1175+
11711176
def __init__(self, n=1, normalize=False, weekday=0):
11721177
BaseOffset.__init__(self, n, normalize)
1173-
object.__setattr__(self, "weekday", weekday)
1178+
self.weekday = weekday
11741179

11751180
if weekday < 0 or weekday > 6:
11761181
raise ValueError(f"Day must be 0<=day<=6, got {weekday}")

pandas/tests/tseries/offsets/test_offsets.py

+18
Original file line numberDiff line numberDiff line change
@@ -655,6 +655,19 @@ def test_pickle_v0_15_2(self, datapath):
655655
#
656656
tm.assert_dict_equal(offsets, read_pickle(pickle_path))
657657

658+
def test_pickle_roundtrip(self, offset_types):
659+
off = self._get_offset(offset_types)
660+
res = tm.round_trip_pickle(off)
661+
assert off == res
662+
if type(off) is not DateOffset:
663+
for attr in off._attributes:
664+
if attr == "calendar":
665+
# np.busdaycalendar __eq__ will return False;
666+
# we check holidays and weekmask attrs so are OK
667+
continue
668+
# Make sure nothings got lost from _params (which __eq__) is based on
669+
assert getattr(off, attr) == getattr(res, attr)
670+
658671
def test_onOffset_deprecated(self, offset_types):
659672
# GH#30340 use idiomatic naming
660673
off = self._get_offset(offset_types)
@@ -3463,6 +3476,11 @@ def test_is_on_offset(self, case):
34633476
offset = LastWeekOfMonth(weekday=weekday)
34643477
assert offset.is_on_offset(dt) == expected
34653478

3479+
def test_repr(self):
3480+
assert (
3481+
repr(LastWeekOfMonth(n=2, weekday=1)) == "<2 * LastWeekOfMonths: weekday=1>"
3482+
)
3483+
34663484

34673485
class TestSemiMonthEnd(Base):
34683486
_offset = SemiMonthEnd

pandas/tseries/offsets.py

+68-1
Original file line numberDiff line numberDiff line change
@@ -310,8 +310,53 @@ def is_on_offset(self, dt):
310310
# TODO, see #1395
311311
return True
312312

313+
def _repr_attrs(self) -> str:
314+
# The DateOffset class differs from other classes in that members
315+
# of self._attributes may not be defined, so we have to use __dict__
316+
# instead.
317+
exclude = {"n", "inc", "normalize"}
318+
attrs = []
319+
for attr in sorted(self.__dict__):
320+
if attr.startswith("_") or attr == "kwds":
321+
continue
322+
elif attr not in exclude:
323+
value = getattr(self, attr)
324+
attrs.append(f"{attr}={value}")
325+
326+
out = ""
327+
if attrs:
328+
out += ": " + ", ".join(attrs)
329+
return out
313330

314-
class BusinessDay(BusinessMixin, SingleConstructorOffset):
331+
@cache_readonly
332+
def _params(self):
333+
"""
334+
Returns a tuple containing all of the attributes needed to evaluate
335+
equality between two DateOffset objects.
336+
"""
337+
# The DateOffset class differs from other classes in that members
338+
# of self._attributes may not be defined, so we have to use __dict__
339+
# instead.
340+
all_paras = self.__dict__.copy()
341+
all_paras["n"] = self.n
342+
all_paras["normalize"] = self.normalize
343+
for key in self.__dict__:
344+
if key not in all_paras:
345+
# cython attributes are not in __dict__
346+
all_paras[key] = getattr(self, key)
347+
348+
if "holidays" in all_paras and not all_paras["holidays"]:
349+
all_paras.pop("holidays")
350+
exclude = ["kwds", "name", "calendar"]
351+
attrs = [
352+
(k, v) for k, v in all_paras.items() if (k not in exclude) and (k[0] != "_")
353+
]
354+
attrs = sorted(set(attrs))
355+
params = tuple([str(type(self))] + attrs)
356+
return params
357+
358+
359+
class BusinessDay(BusinessMixin):
315360
"""
316361
DateOffset subclass representing possibly n business days.
317362
"""
@@ -796,6 +841,22 @@ def __init__(
796841
BusinessHour.__init__(self, n, normalize, start=start, end=end, offset=offset)
797842
CustomMixin.__init__(self, weekmask, holidays, calendar)
798843

844+
def __reduce__(self):
845+
# None for self.calendar bc np.busdaycalendar doesnt pickle nicely
846+
return (
847+
type(self),
848+
(
849+
self.n,
850+
self.normalize,
851+
self.weekmask,
852+
self.holidays,
853+
None,
854+
self.start,
855+
self.end,
856+
self.offset,
857+
),
858+
)
859+
799860

800861
# ---------------------------------------------------------------------
801862
# Month-Based Offset Classes
@@ -1286,6 +1347,9 @@ def __init__(self, n=1, normalize=False, week=0, weekday=0):
12861347
if self.week < 0 or self.week > 3:
12871348
raise ValueError(f"Week must be 0<=week<=3, got {self.week}")
12881349

1350+
def __reduce__(self):
1351+
return type(self), (self.n, self.normalize, self.week, self.weekday)
1352+
12891353
def _get_offset_day(self, other: datetime) -> int:
12901354
"""
12911355
Find the day in the same month as other that has the same
@@ -1345,6 +1409,9 @@ def __init__(self, n=1, normalize=False, weekday=0):
13451409
raise ValueError("N cannot be 0")
13461410
object.__setattr__(self, "week", -1)
13471411

1412+
def __reduce__(self):
1413+
return type(self), (self.n, self.normalize, self.weekday)
1414+
13481415
def _get_offset_day(self, other: datetime) -> int:
13491416
"""
13501417
Find the day in the same month as other that has the same

0 commit comments

Comments
 (0)