debian-workalendar/workalendar/core.py

133 lines
3.9 KiB
Python

"""Workday tools
"""
from calendar import monthrange
from datetime import date, timedelta
from dateutil import easter
MON, TUE, WED, THU, FRI, SAT, SUN = range(7)
class Calendar(object):
EASTER_METHOD = 3 # 3 is 'Western'
def __init__(self):
self._holidays = {}
def get_calendar_holidays(self, year):
"""Get calendar holidays.
This method **must** return a set or a list.
You must override this method for each calendar."""
return set([])
def holidays(self, year=None):
"Computes holidays (non-working days) for a given year"
if not year:
year = date.today().year
if year in self._holidays:
return self._holidays[year]
if year not in self._holidays:
self._holidays[year] = set([])
# Here we process the holiday specific calendar
self._holidays[year] = self.get_calendar_holidays(year)
return set(self._holidays[year])
def get_weekend_days(self):
"""Return a list (or a tuple) of weekdays that are *not* workdays.
e.g: return (SAT, SUN,)
"""
raise NotImplementedError("Your Calendar class must implement the"
" `get_weekend_days` method")
def is_workday(self, day):
"Return True if it's a workday."
if day.weekday() in self.get_weekend_days():
return False
if day in self.holidays(day.year):
return False
return True
def add_workdays(self, day, delta):
"Add `delta` workdays to the date."
days = 0
temp_day = day
while days < delta:
temp_day = temp_day + timedelta(1)
if self.is_workday(temp_day):
days += 1
return temp_day
def get_easter_sunday(self, year):
"Return the date of the easter (sunday) -- following the easter method"
return easter.easter(year, self.EASTER_METHOD)
def get_easter_monday(self, year):
"Return the date of the monday after easter"
sunday = self.get_easter_sunday(year)
return sunday + timedelta(days=1)
@staticmethod
def get_nth_weekday_in_month(year, month, weekday, n=1):
"""Get the nth weekday in a given month. e.g:
>>> # the 1st monday in Jan 2013
>>> Calendar.get_nth_weekday_in_month(2013, 1, MON)
datetime.date(2013, 1, 7)
>>> # The 2nd monday in Jan 2013
>>> Calendar.get_nth_weekday_in_month(2013, 1, MON, 2)
datetime.date(2013, 1, 14)
"""
day = date(year, month, 1)
counter = 0
while True:
if day.month != month:
# Don't forget to break if "n" is too big
return None
if day.weekday() == weekday:
counter += 1
if counter == n:
break
day = day + timedelta(days=1)
return day
@staticmethod
def get_last_weekday_in_month(year, month, weekday):
"""Get the last weekday in a given month. e.g:
>>> # the last monday in Jan 2013
>>> Calendar.get_last_weekday_in_month(2013, 1, MON)
datetime.date(2013, 1, 28)
"""
day = date(year, month, monthrange(year, month)[1])
while True:
if day.weekday() == weekday:
break
day = day - timedelta(days=1)
return day
class WesternCalendar(Calendar):
"""
General usage calendar for Western countries.
(chiefly Europe and Northern America)
"""
EASTER_METHOD = 3 # 3 is 'Western'
WEEK_END_DAYS = (SAT, SUN)
def get_calendar_holidays(self, year):
"European countries have at least these 2 days as holidays in common"
days = set([])
days.add(date(year, 1, 1))
days.add(date(year, 12, 25))
return days
def get_weekend_days(self):
"Week-end days are SATurday and SUNday."
return self.WEEK_END_DAYS