휴일 표현과 계산.

days.py

import datetime


class DayExpression(object):
    def include(self, date):
        raise NotImplementedError


class Union(DayExpression):
    def __init__(self, expressions):
        self.expressions = expressions

    def include(self, date):
        for e in self.expressions:
            if e.include(date):
                return True
        return False


class DayOfWeek(DayExpression):
    def __init__(self, day):
        self.day = day

    def include(self, date):
        return date.weekday() == self.day

    def __eq__(self, other):
        return self.day == other.day

for name, value in zip(["MON", "TUE", "WED", "THU", "FRI", "SAT", "SUN"], range(0, 7)):
    setattr(DayOfWeek, name, DayOfWeek(value))


class MonthDay(DayExpression):
    def __init__(self, month, day):
        self.month = month
        self.day = day

    def include(self, date):
        return self.month == date.month and self.day == date.day

    def __eq__(self, other):
        return self.month == other.month and self.day == other.day


class Date(DayExpression):
    def __init__(self, year, month, day):
        self.year = year
        self.month = month
        self.day = day

    def include(self, date):
        return self.year == date.year and self.month == date.month and self.day == date.day

    def __eq__(self, other):
        return self.year == other.year and self.month == other.month and self.day == other.day


class Holyday(Union):
    def __init__(self, name, expressions):
        super(Holyday, self).__init__(expressions)
        self.name = name

    def __eq__(self, other):
        return self.name == other.name and self.expressions == other.expressions


class HolydayCalendar(Union):
    def __init__(self, holydays):
        super(HolydayCalendar, self).__init__(holydays)

    def __eq__(self, other):
        return self.expressions == other.expressions

    @property
    def holydays(self):
        return self.expressions

    def workday(self, base, offset):
        day = base
        while True:
            if not self.include(day):
                if offset == 0:
                    return day
                offset -= 1
            day = day + datetime.timedelta(days=1)

    @classmethod
    def load_from_json(cls, json):
        holydays = []
        for name, tuples in json:
            exprs = []
            for tuple in tuples:
                if tuple[0] == 'DayOfWeek':
                    expr = DayOfWeek({"MON": 0, "TUE": 1, "WED": 2, "THU": 3, "FRI": 4, "SAT": 5, "SUN": 6}[tuple[1].upper()])
                elif tuple[0] == 'MonthDay':
                    expr = MonthDay(tuple[1], tuple[2])
                elif tuple[0] == 'Date':
                    expr = Date(tuple[1], tuple[2], tuple[3])
                else:
                    raise Exception("Unknown expression: %s" % tuple[0])
                exprs.append(expr)
            holydays.append(Holyday(name, exprs))
        return HolydayCalendar(holydays)

    @classmethod
    def dump_to_json(cls, calendar):
        json = []
        for holyday in calendar.expressions:
            tuples = []
            for expr in holyday.expressions:
                if isinstance(expr, DayOfWeek):
                    tuple = ('DayOfWeek', ["MON", "TUE", "WED", "THU", "FRI", "SAT", "SUN"][expr.day])
                elif isinstance(expr, MonthDay):
                    tuple = ('MonthDay', expr.month, expr.day)
                elif isinstance(expr, Date):
                    tuple = ('Date', expr.year, expr.month, expr.day)
                else:
                    raise Exception("Unknown expression: %s" % expr.__class__.__name__)
                tuples.append(tuple)
            json.append((holyday.name, tuples))
        return json

days_test.py

import days
import datetime
import unittest


class DaysExpressionTest(unittest.TestCase):
    def test_day_of_week(self):
        dut = days.DayOfWeek.SUN
        sat = datetime.date(2013, 10, 12)
        sun = datetime.date(2013, 10, 13)
        self.assertTrue(dut.include(sun))
        self.assertFalse(dut.include(sat))

    def test_month_day(self):
        dut = days.MonthDay(10, 3)
        self.assertTrue(dut.include(datetime.date(2013, 10, 3)))
        self.assertTrue(dut.include(datetime.date(2012, 10, 3)))
        self.assertFalse(dut.include(datetime.date(2012, 10, 4)))

    def test_date(self):
        dut = days.Date(2013, 10, 10)
        self.assertTrue(dut.include(datetime.date(2013, 10, 10)))
        self.assertFalse(dut.include(datetime.date(2013, 10, 11)))

    def test_union(self):
        dut = days.Union([days.DayOfWeek.SUN, days.Date(2013, 10, 3)])
        self.assertTrue(dut.include(datetime.date(2013, 10, 3)) and dut.include(datetime.date(2013, 10, 6)))
        self.assertFalse(dut.include(datetime.date(2013, 10, 4)))


class HolydayCalendarTest(unittest.TestCase):
    def test_workday_0th(self):
        dut = days.HolydayCalendar([days.Holyday("Sunday", [days.DayOfWeek.SUN])])
        self.assertEquals(datetime.date(2013, 10, 12), dut.workday(datetime.date(2013, 10, 12), 0))
        self.assertEquals(datetime.date(2013, 10, 14), dut.workday(datetime.date(2013, 10, 13), 0))
        self.assertEquals(datetime.date(2013, 10, 14), dut.workday(datetime.date(2013, 10, 14), 0))

    def test_workday_1st(self):
        dut = days.HolydayCalendar([days.Holyday("Sunday", [days.DayOfWeek.SUN])])
        self.assertEquals(datetime.date(2013, 10, 14), dut.workday(datetime.date(2013, 10, 12), 1))
        self.assertEquals(datetime.date(2013, 10, 15), dut.workday(datetime.date(2013, 10, 13), 1))
        self.assertEquals(datetime.date(2013, 10, 15), dut.workday(datetime.date(2013, 10, 14), 1))

    def test_dump_empty(self):
        self.assertEqual(days.HolydayCalendar.dump_to_json(days.HolydayCalendar([])), [])

    def test_load_empty(self):
        self.assertEqual(days.HolydayCalendar.load_from_json([]), days.HolydayCalendar([]))

    def test_dump(self):
        self.assertEqual(days.HolydayCalendar.dump_to_json(days.HolydayCalendar([days.Holyday("Sunday", [days.DayOfWeek.SUN])])),
                         [("Sunday", [("DayOfWeek", "SUN")])])

    def test_load(self):
        self.assertEqual(days.HolydayCalendar.load_from_json([("Sunday", [("DayOfWeek", "SUN")])]),
                         days.HolydayCalendar([days.Holyday("Sunday", [days.DayOfWeek.SUN])]))