Skip to content

Commit 1517e1b

Browse files
jseaboldadamklein
authored andcommitted
ENH: Add BMonthBegin and BQuarterBegin
1 parent 1ca8bc7 commit 1517e1b

File tree

1 file changed

+111
-42
lines changed

1 file changed

+111
-42
lines changed

pandas/core/datetools.py

Lines changed: 111 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,19 @@ def normalize_date(dt):
102102
dt = _dt_box(dt)
103103
return datetime(dt.year, dt.month, dt.day)
104104

105+
def _get_firstbday(wkday):
106+
"""
107+
wkday is the result of calendar.monthrange(year, month)
108+
109+
If it's a saturday or sunday, increment first business day to reflect this
110+
"""
111+
firstBDay = 1
112+
if wkday == 5: # on Saturday
113+
firstBDay = 3
114+
elif wkday == 6: # on Sunday
115+
firstBDay = 2
116+
return firstBDay
117+
105118
#-------------------------------------------------------------------------------
106119
# DateOffset
107120

@@ -372,7 +385,7 @@ def onOffset(cls, someDate):
372385
return someDate.day == days_in_month
373386

374387
class MonthBegin(DateOffset, CacheableOffset):
375-
"""DateOFfset of one month begin"""
388+
"""DateOffset of one month at beginning"""
376389

377390
_normalizeFirst = True
378391

@@ -415,6 +428,28 @@ def apply(self, other):
415428
other = other - BDay()
416429
return other
417430

431+
class BMonthBegin(DateOffset, CacheableOffset):
432+
"""DateOffset of one business month at beginning"""
433+
434+
_normalizeFirst = True
435+
436+
def apply(self, other):
437+
n = self.n
438+
439+
wkday, _ = calendar.monthrange(other.year, other.month)
440+
firstBDay = _get_firstbday(wkday)
441+
442+
if other.day > firstBDay:
443+
if n <= 0:
444+
# as if rolled forward already
445+
n = n + 1
446+
447+
other = other + lib.Delta(months=n)
448+
wkday, _ = calendar.monthrange(other.year, other.month)
449+
firstBDay = _get_firstbday(wkday)
450+
result = datetime(other.year, other.month, firstBDay)
451+
return result
452+
418453

419454
class Week(DateOffset, CacheableOffset):
420455
"""
@@ -581,6 +616,49 @@ def onOffset(self, someDate):
581616
modMonth = (someDate.month - self.startingMonth) % 3
582617
return BMonthEnd().onOffset(someDate) and modMonth == 0
583618

619+
class BQuarterBegin(DateOffset, CacheableOffset):
620+
_outputName = "BusinessQuarterBegin"
621+
_normalizeFirst = True
622+
623+
def __init__(self, n=1, **kwds):
624+
self.n = n
625+
self.startingMonth = kwds.get('startingMonth', 3)
626+
627+
self.offset = BMonthBegin(3)
628+
self.kwds = kwds
629+
630+
def isAnchored(self):
631+
return (self.n == 1 and self.startingMonth is not None)
632+
633+
def apply(self, other):
634+
n = self.n
635+
636+
if self._normalizeFirst:
637+
other = normalize_date(other)
638+
639+
wkday, _ = calendar.monthrange(other.year, other.month)
640+
641+
firstBDay = _get_firstbday(wkday)
642+
643+
monthsSince = (other.month - self.startingMonth) % 3
644+
if monthsSince == 3: # on offset
645+
monthsSince = 0
646+
647+
if n <= 0 and monthsSince != 0: # make sure to roll forward so negate
648+
monthsSince = monthsSince - 3
649+
if n <= 0 and (monthsSince == 0 and other.day > firstBDay):
650+
n = n + 1
651+
elif n > 0 and (monthsSince == 0 and other.day < firstBDay):
652+
n = n - 1
653+
654+
655+
other = other + lib.Delta(months=3*n - monthsSince)
656+
wkday, _ = calendar.monthrange(other.year, other.month)
657+
firstBDay = _get_firstbday(wkday)
658+
result = datetime(other.year, other.month, firstBDay)
659+
return result
660+
661+
584662
class QuarterEnd(DateOffset, CacheableOffset):
585663
"""DateOffset increments between business Quarter dates
586664
startingMonth = 1 corresponds to dates like 1/31/2007, 4/30/2007, ...
@@ -718,11 +796,7 @@ def apply(self, other):
718796

719797
wkday, days_in_month = calendar.monthrange(other.year, self.month)
720798

721-
firstBDay = 1
722-
if wkday == 5: # on Saturday
723-
firstBDay = 3
724-
elif wkday == 6: # on Sunday
725-
firstBDay = 2
799+
firstBDay = _get_firstbday(wkday)
726800

727801
years = n
728802
if n > 0:
@@ -738,12 +812,7 @@ def apply(self, other):
738812

739813
wkday, days_in_month = calendar.monthrange(other.year, self.month)
740814

741-
742-
firstBDay = 1
743-
if wkday == 5: # on Saturday
744-
firstBDay = 3
745-
elif wkday == 6: # on Sunday
746-
firstBDay = 2
815+
firstBDay = _get_firstbday(wkday)
747816

748817
result = datetime(other.year, self.month, firstBDay)
749818
return result
@@ -948,19 +1017,19 @@ class Second(Tick):
9481017
"Q@NOV" : QuarterEnd(startingMonth=11),
9491018
"Q@DEC" : QuarterEnd(startingMonth=12),
9501019
# Quarterly - Calendar (Start)
951-
#"QS@JAN" : QuarterBegin(startingMonth=1),
952-
#"QS" : QuarterBegin(startingMonth=1),
953-
#"QS@FEB" : QuarterBegin(startingMonth=2),
954-
#"QS@MAR" : QuarterBegin(startingMonth=3),
955-
#"QS@APR" : QuarterBegin(startingMonth=4),
956-
#"QS@MAY" : QuarterBegin(startingMonth=5),
957-
#"QS@JUN" : QuarterBegin(startingMonth=6),
958-
#"QS@JUL" : QuarterBegin(startingMonth=7),
959-
#"QS@AUG" : QuarterBegin(startingMonth=8),
960-
#"QS@SEP" : QuarterBegin(startingMonth=9),
961-
#"QS@OCT" : QuarterBegin(startingMonth=10),
962-
#"QS@NOV" : QuarterBegin(startingMonth=11),
963-
#"QS@DEC" : QuarterBegin(startingMonth=12),
1020+
"QS@JAN" : QuarterBegin(startingMonth=1),
1021+
"QS" : QuarterBegin(startingMonth=1),
1022+
"QS@FEB" : QuarterBegin(startingMonth=2),
1023+
"QS@MAR" : QuarterBegin(startingMonth=3),
1024+
"QS@APR" : QuarterBegin(startingMonth=4),
1025+
"QS@MAY" : QuarterBegin(startingMonth=5),
1026+
"QS@JUN" : QuarterBegin(startingMonth=6),
1027+
"QS@JUL" : QuarterBegin(startingMonth=7),
1028+
"QS@AUG" : QuarterBegin(startingMonth=8),
1029+
"QS@SEP" : QuarterBegin(startingMonth=9),
1030+
"QS@OCT" : QuarterBegin(startingMonth=10),
1031+
"QS@NOV" : QuarterBegin(startingMonth=11),
1032+
"QS@DEC" : QuarterBegin(startingMonth=12),
9641033
# Quarterly - Business
9651034
"BQ@JAN" : BQuarterEnd(startingMonth=1),
9661035
"BQ@FEB" : BQuarterEnd(startingMonth=2),
@@ -976,29 +1045,29 @@ class Second(Tick):
9761045
"BQ@NOV" : BQuarterEnd(startingMonth=11),
9771046
"BQ@DEC" : BQuarterEnd(startingMonth=12),
9781047
# Quarterly - Business (Start)
979-
#"BQS@JAN" : BQuarterBegin(startingMonth=1),
980-
#"BQS" : BQuarterBegin(startingMonth=1),
981-
#"BQS@FEB" : BQuarterBegin(startingMonth=2),
982-
#"BQS@MAR" : BQuarterBegin(startingMonth=3),
983-
#"BQS@APR" : BQuarterBegin(startingMonth=4),
984-
#"BQS@MAY" : BQuarterBegin(startingMonth=5),
985-
#"BQS@JUN" : BQuarterBegin(startingMonth=6),
986-
#"BQS@JUL" : BQuarterBegin(startingMonth=7),
987-
#"BQS@AUG" : BQuarterBegin(startingMonth=8),
988-
#"BQS@SEP" : BQuarterBegin(startingMonth=9),
989-
#"BQS@OCT" : BQuarterBegin(startingMonth=10),
990-
#"BQS@NOV" : BQuarterBegin(startingMonth=11),
991-
#"BQS@DEC" : BQuarterBegin(startingMonth=12),
1048+
"BQS@JAN" : BQuarterBegin(startingMonth=1),
1049+
"BQS" : BQuarterBegin(startingMonth=1),
1050+
"BQS@FEB" : BQuarterBegin(startingMonth=2),
1051+
"BQS@MAR" : BQuarterBegin(startingMonth=3),
1052+
"BQS@APR" : BQuarterBegin(startingMonth=4),
1053+
"BQS@MAY" : BQuarterBegin(startingMonth=5),
1054+
"BQS@JUN" : BQuarterBegin(startingMonth=6),
1055+
"BQS@JUL" : BQuarterBegin(startingMonth=7),
1056+
"BQS@AUG" : BQuarterBegin(startingMonth=8),
1057+
"BQS@SEP" : BQuarterBegin(startingMonth=9),
1058+
"BQS@OCT" : BQuarterBegin(startingMonth=10),
1059+
"BQS@NOV" : BQuarterBegin(startingMonth=11),
1060+
"BQS@DEC" : BQuarterBegin(startingMonth=12),
9921061
# Monthly - Calendar
9931062
"M" : MonthEnd(),
9941063
"EOM" : MonthEnd(),
995-
#"MS" : MonthBegin(),
996-
#"SOM" : MonthBegin(),
1064+
"MS" : MonthBegin(),
1065+
"SOM" : MonthBegin(),
9971066
# Monthly - Business
9981067
"BM" : BMonthEnd(),
9991068
"BEOM" : BMonthEnd(),
1000-
#"BMS" : BMonthBegin(),
1001-
#"BSOM" : BMonthBegin(),
1069+
"BMS" : BMonthBegin(),
1070+
"BSOM" : BMonthBegin(),
10021071
# Weekly
10031072
"W@MON" : Week(weekday=0),
10041073
"WS" : Week(weekday=0),

0 commit comments

Comments
 (0)