class BaseSchedule(object):
__metaclass__ = ScheduleMetaClass
calendar = None
open_time = MIN_TIME
close_time = MAX_TIME
tzinfo = pytz.UTC
weekmask = 'Mon Tue Wed Thu Fri'
def __init__(self):
self.cbd = CustomBusinessDay(calendar=self.calendar,
weekmask=self.weekmask)
self.delta = timedelta_between(self.close_time, self.open_time)
def get_schedule_on(self, date):
start = to_tz_datetime(date=date,
time=self.open_time,
to_tz=self.tzinfo)
end = start + self.delta
if self.cbd.onOffset(start) and self.cbd.onOffset(end):
return start, end
class CustomBusinessYearEnd(YearOffset):
_cacheable = False
_prefix = 'CBYE'
_default_month = 12
def __init__(self,
n=1,
normalize=False,
weekmask='Mon Tue Wed Thu Fri',
holidays=None,
calendar=None,
**kwds):
self.n = n
self.normalize = normalize
self.kwds.update(kwds)
self.offset = kwds.get('offset', timedelta(0))
self.month = kwds.get('month', self._default_month)
try:
kwds.pop('month')
except Exception as e:
pass
self.cbday = CustomBusinessDay(n=1,
normalize=normalize,
weekmask=weekmask,
holidays=holidays,
calendar=calendar,
**kwds)
self.kwds['calendar'] = self.cbday.calendar
self.y_offset = YearEnd(1)
@apply_wraps
def apply(self, other):
n = self.n
cur_yend = self.y_offset.rollforward(other)
cur_cyend = self.cbday.rollback(cur_yend)
if n == 0 and other != cur_cyend:
n += 1
if other < cur_cyend and n >= 1:
n -= 1
if other > cur_cyend and n <= -1:
n += 1
new = cur_yend + n * self.y_offset
result = self.cbday.rollback(new)
return result
def onOffset(self, dt):
if self.normalize and not _is_normalized(dt):
return False
if not self.cbday.onOffset(dt):
return False
return (dt + self.cbday).year != dt.year
class CustomBusinessQuaterEnd(QuarterOffset):
_cacheable = False
_prefix = 'CBQE'
_attributes = frozenset({'holidays', 'calendar'}
| set(QuarterOffset._attributes))
def __init__(self,
n=1,
normalize=False,
weekmask='Mon Tue Wed Thu Fri',
holidays=None,
calendar=None,
**kwds):
self.n = n
self.normalize = normalize
self.kwds.update(kwds)
self.offset = kwds.get('offset', timedelta(0))
self.startingMonth = kwds.get('startingMonth', 3)
self.cbday = CustomBusinessDay(n=1,
normalize=normalize,
weekmask=weekmask,
holidays=holidays,
calendar=calendar)
self.calendar = self.cbday.calendar
self.holidays = holidays
self.startingMonth = self.startingMonth
self.q_offset = QuarterEnd(1)
@apply_wraps
def apply(self, other):
n = self.n
cur_qend = self.q_offset.rollforward(other)
cur_cqend = self.cbday.rollback(cur_qend)
if n == 0 and other != cur_cqend:
n += 1
if other < cur_cqend and n >= 1:
n -= 1
if other > cur_cqend and n <= -1:
n += 1
new = cur_qend + n * self.q_offset
result = self.cbday.rollback(new)
return result
def onOffset(self, dt):
if self.normalize and not _is_normalized(dt):
return False
if not self.cbday.onOffset(dt):
return False
return (dt + self.cbday).quarter != dt.quarter
class CustomBusinessWeekEnd(DateOffset):
_cacheable = False
_prefix = 'CBWE'
_attributes = frozenset(['calendar', 'holidays'])
def __init__(self,
n=1,
normalize=False,
weekmask='Mon Tue Wed Thu Fri',
holidays=None,
calendar=None,
**kwds):
self.n = n
self.normalized = normalize
self.kwds.update(kwds)
self.offset = kwds.get('offset', timedelta(0))
self.cbday = CustomBusinessDay(n=1,
normalize=normalize,
weekmask=weekmask,
holidays=holidays,
calendar=calendar,
offset=self.offset)
self.calendar = self.cbday.calendar
self.holidays = holidays
self.w_offset = Week(weekday=4)
@apply_wraps
def apply(self, other):
n = self.n
result = other
if n == 0:
n = 1
if n > 0:
while result <= other:
next_fri = other + n * self.w_offset
result = self.cbday.rollback(next_fri)
n += 1
else:
while result >= other:
last_fri = other + n * self.w_offset
result = self.cbday.rollback(last_fri)
n -= 1
return result
def onOffset(self, dt):
if self.normalize and not _is_normalized(dt):
return False
if not self.cbday.onOffset(dt):
return False
return (dt + self.cbday).week != dt.week
