lgc.date should now be usable with either mx or pydatetime objects

2 files changed, 223 insertions, 149 deletions

diff --git a/date.py b/date.py
index a49b2db..bcdb0f9 100644
--- a/date.py
+++ b/date.py
@@ -1,6 +1,6 @@
 """Date manipulation helper functions.
 
-:copyright: 2006-2008 LOGILAB S.A. (Paris, FRANCE), all rights reserved.
+:copyright: 2006-2009 LOGILAB S.A. (Paris, FRANCE), all rights reserved.
 :contact: http://www.logilab.fr/ -- mailto:contact@logilab.fr
 :license: General Public License version 2 - http://www.gnu.org/licenses
 """
@@ -8,100 +8,159 @@ __docformat__ = "restructuredtext en"
 
 import math
 
+from datetime import date, datetime, timedelta
 try:
-    from mx.DateTime import RelativeDateTime, strptime, Date
-    STEP = 1
+    from mx.DateTime import RelativeDateTime, Date
 except ImportError:
-    from warnings import warn
-    warn("mxDateTime not found, holiday management won't be available")
-    from datetime import timedelta
-    STEP = timedelta(days=1)
+    from datetime import date, timedelta
+    warn("mxDateTime not found, endsOfMonth won't be available")
+    def weekday(date):
+        return date.weekday()
+    endOfMonth = None
 else:
     endOfMonth = RelativeDateTime(months=1, day=-1)
 
-    FRENCH_FIXED_HOLIDAYS = {
-        'jour_an'        : '%s-01-01',
-        'fete_travail'   : '%s-05-01',
-        'armistice1945'  : '%s-05-08',
-        'fete_nat'       : '%s-07-14',
-        'assomption'     : '%s-08-15',
-        'toussaint'      : '%s-11-01',
-        'armistice1918'  : '%s-11-11',
-        'noel'           : '%s-12-25',
-        }
-
-
-    FRENCH_MOBILE_HOLIDAYS = {
-        'paques2004'    : '2004-04-12',
-        'ascension2004' : '2004-05-20',
-        'pentecote2004' : '2004-05-31',
-
-        'paques2005'    : '2005-03-28',
-        'ascension2005' : '2005-05-05',
-        'pentecote2005' : '2005-05-16',
-
-        'paques2006'    : '2006-04-17',
-        'ascension2006' : '2006-05-25',
-        'pentecote2006' : '2006-06-05',
-
-        'paques2007'    : '2007-04-09',
-        'ascension2007' : '2007-05-17',
-        'pentecote2007' : '2007-05-28',
-
-        'paques2008'    : '2008-03-24',
-        'ascension2008' : '2008-05-01',
-        'pentecote2008' : '2008-05-12',
-        }
-
-    def get_national_holidays(begin, end):
-        """return french national days off between begin and end"""
-        begin = Date(begin.year, begin.month, begin.day)
-        end = Date(end.year, end.month, end.day)
-        holidays = [strptime(datestr, '%Y-%m-%d')
-                    for datestr in FRENCH_MOBILE_HOLIDAYS.values()]
-        for year in xrange(begin.year, end.year+1):
-            for datestr in FRENCH_FIXED_HOLIDAYS.values():
-                date = strptime(datestr % year, '%Y-%m-%d')
-                if date not in holidays:
-                    holidays.append(date)
-        return [day for day in holidays if begin <= day < end]
-
-
-    def add_days_worked(start, days):
-        """adds date but try to only take days worked into account"""
-        weeks, plus = divmod(days, 5)
-        end = start+(weeks * 7) + plus
-        if end.day_of_week >= 5: # saturday or sunday
-            end += 2
-        end += len([x for x in get_national_holidays(start, end+1)
-                    if x.day_of_week < 5])
-        if end.day_of_week >= 5: # saturday or sunday
-            end += 2
-        return end
-
-    def nb_open_days(start, end):
-        assert start <= end
-        days = int(math.ceil((end - start).days))
-        weeks, plus = divmod(days, 7)
-        if start.day_of_week > end.day_of_week:
-            plus -= 2
-        elif end.day_of_week == 6:
-            plus -= 1
-        open_days = weeks * 5 + plus
-        nb_week_holidays = len([x for x in get_national_holidays(start, end+1)
-                                if x.day_of_week < 5 and x < end])
-        return open_days  - nb_week_holidays
-
-
-def date_range(begin, end, step=STEP):
+# NOTE: should we implement a compatibility layer between date representations
+#       as we have in lgc.db ?
+
+PYDATE_STEP = timedelta(days=1)
+
+FRENCH_FIXED_HOLIDAYS = {
+    'jour_an'        : '%s-01-01',
+    'fete_travail'   : '%s-05-01',
+    'armistice1945'  : '%s-05-08',
+    'fete_nat'       : '%s-07-14',
+    'assomption'     : '%s-08-15',
+    'toussaint'      : '%s-11-01',
+    'armistice1918'  : '%s-11-11',
+    'noel'           : '%s-12-25',
+    }
+
+FRENCH_MOBILE_HOLIDAYS = {
+    'paques2004'    : '2004-04-12',
+    'ascension2004' : '2004-05-20',
+    'pentecote2004' : '2004-05-31',
+
+    'paques2005'    : '2005-03-28',
+    'ascension2005' : '2005-05-05',
+    'pentecote2005' : '2005-05-16',
+
+    'paques2006'    : '2006-04-17',
+    'ascension2006' : '2006-05-25',
+    'pentecote2006' : '2006-06-05',
+
+    'paques2007'    : '2007-04-09',
+    'ascension2007' : '2007-05-17',
+    'pentecote2007' : '2007-05-28',
+
+    'paques2008'    : '2008-03-24',
+    'ascension2008' : '2008-05-01',
+    'pentecote2008' : '2008-05-12',
+
+    'paques2009'    : '2009-04-13',
+    'ascension2009' : '2009-05-21',
+    'pentecote2009' : '2009-06-01',
+
+    'paques2010'    : '2010-04-05',
+    'ascension2010' : '2010-05-13',
+    'pentecote2010' : '2010-05-24',
+
+    'paques2011'    : '2011-04-25',
+    'ascension2011' : '2011-06-02',
+    'pentecote2011' : '2011-06-13',
+
+    'paques2012'    : '2012-04-09',
+    'ascension2012' : '2012-05-17',
+    'pentecote2012' : '2012-05-28',
+    }
+
+# this implementation cries for multimethod dispatching
+
+def get_step(dateobj):
+    # assume date is either a python datetime or a mx.DateTime object
+    if isinstance(dateobj, date):
+        return PYDATE_STEP
+    return 1 # mx.DateTime is ok with integers
+
+def datefactory(year, month, day, sampledate):
+    # assume date is either a python datetime or a mx.DateTime object
+    if isinstance(sampledate, datetime):
+        return datetime(year, month, day)
+    if isinstance(sampledate, date):
+        return date(year, month, day)
+    return Date(year, month, day)
+
+def weekday(dateobj):
+    # assume date is either a python datetime or a mx.DateTime object
+    if isinstance(dateobj, date):
+        return dateobj.weekday()
+    return dateobj.day_of_week
+
+def str2date(datestr, sampledate):
+    # NOTE: datetime.strptime is not an option until we drop py2.4 compat
+    year, month, day = [int(chunk) for chunk in datestr.split('-')]
+    return datefactory(year, month, day, sampledate)
+
+def days_between(start, end):
+    if isinstance(start, date):
+        delta = end - start
+        # datetime.timedelta.days is always an integer (floored)
+        if delta.seconds:
+            return delta.days + 1
+        return delta.days
+    else:
+        return int(math.ceil((end - start).days))
+
+def get_national_holidays(begin, end):
+    """return french national days off between begin and end"""
+    begin = datefactory(begin.year, begin.month, begin.day, begin)
+    end = datefactory(end.year, end.month, end.day, end)
+    holidays = [str2date(datestr, begin)
+                for datestr in FRENCH_MOBILE_HOLIDAYS.values()]
+    for year in xrange(begin.year, end.year+1):
+        for datestr in FRENCH_FIXED_HOLIDAYS.values():
+            date = str2date(datestr % year, begin)
+            if date not in holidays:
+                holidays.append(date)
+    return [day for day in holidays if begin <= day < end]
+
+def add_days_worked(start, days):
+    """adds date but try to only take days worked into account"""
+    step = get_step(start)
+    weeks, plus = divmod(days, 5)
+    end = start + ((weeks * 7) + plus) * step
+    if weekday(end) >= 5: # saturday or sunday
+        end += (2 * step)
+    end += len([x for x in get_national_holidays(start, end + step)
+                if weekday(x) < 5]) * step
+    if weekday(end) >= 5: # saturday or sunday
+        end += (2 * step)
+    return end
+
+def nb_open_days(start, end):
+    assert start <= end
+    step = get_step(start)
+    days = days_between(start, end)
+    weeks, plus = divmod(days, 7)
+    if weekday(start) > weekday(end):
+        plus -= 2
+    elif weekday(end) == 6:
+        plus -= 1
+    open_days = weeks * 5 + plus
+    nb_week_holidays = len([x for x in get_national_holidays(start, end+step)
+                            if weekday(x) < 5 and x < end])
+    return open_days - nb_week_holidays
+
+def date_range(begin, end, step=None):
     """
     enumerate dates between begin and end dates.
 
     step can either be oneDay, oneHour, oneMinute, oneSecond, oneWeek
     use endOfMonth to enumerate months
     """
+    if step is None:
+        step = get_step(begin)
     date = begin
     while date < end :
         yield date
         date += step
-
diff --git a/test/unittest_date.py b/test/unittest_date.py
index 6ba98a9..975db7b 100644
--- a/test/unittest_date.py
+++ b/test/unittest_date.py
@@ -1,127 +1,142 @@
 """
 Unittests for date helpers
 """
-
 from logilab.common.testlib import TestCase, unittest_main
 
-from logilab.common.date import date_range
+from logilab.common.date import date_range, endOfMonth
+
+from datetime import date, datetime, timedelta
 
 try:
-    from mx.DateTime import Date, RelativeDate, RelativeDateTime, now, DateTime
-    from logilab.common.date import endOfMonth, add_days_worked, nb_open_days, \
+    from mx.DateTime import Date as mxDate, DateTime as mxDateTime, \
+         now as mxNow, RelativeDateTime, RelativeDate
+    from logilab.common.date import add_days_worked, nb_open_days, \
          get_national_holidays
 except ImportError:
-    from datetime import date as Date, datetime as DateTime, timedelta as RelativeDateTime
-    now = DateTime.now
-    get_national_holidays = endOfMonth = add_days_worked = nb_open_days = None
+    mxDate = mxDateTime = RelativeDateTime = mxNow = None
 
 class DateTC(TestCase):
-    
+    datecls = date
+    datetimecls = datetime
+    timedeltacls = timedelta
+    now = datetime.now
+
     def test_day(self):
         """enumerate days"""
-        r = list(date_range(Date(2000,1,1), Date(2000,1,4)))
-        expected = [Date(2000,1,1), Date(2000,1,2), Date(2000,1,3)]
+        r = list(date_range(self.datecls(2000,1,1), self.datecls(2000,1,4)))
+        expected = [self.datecls(2000,1,1), self.datecls(2000,1,2), self.datecls(2000,1,3)]
         self.assertListEquals(r, expected)
-        r = list(date_range(Date(2000,1,31), Date(2000,2,3)))
-        expected = [Date(2000,1,31), Date(2000,2,1), Date(2000,2,2)]
+        r = list(date_range(self.datecls(2000,1,31), self.datecls(2000,2,3)))
+        expected = [self.datecls(2000,1,31), self.datecls(2000,2,1), self.datecls(2000,2,2)]
         self.assertListEquals(r, expected)
 
-    def test_month(self):
-        """enumerate months"""
-        self.check_mx()
-        r = list(date_range(Date(2000,1,2), Date(2000,4,4), endOfMonth))
-        expected = [Date(2000,1,2), Date(2000,2,29), Date(2000,3,31)]
-        self.assertListEquals(r, expected)
-        r = list(date_range(Date(2000,11,30), Date(2001,2,3), endOfMonth))
-        expected = [Date(2000,11,30), Date(2000,12,31), Date(2001,1,31)]
-        self.assertListEquals(r, expected)
+    def test_add_days_worked_with_mx(self):
+        add = add_days_worked
+        # normal
+        self.assertEquals(add(self.datecls(2008, 1, 3), 1), self.datecls(2008, 1, 4))
+        # skip week-end
+        self.assertEquals(add(self.datecls(2008, 1, 3), 2), self.datecls(2008, 1, 7))
+        # skip 2 week-ends
+        self.assertEquals(add(self.datecls(2008, 1, 3), 8), self.datecls(2008, 1, 15))
+        # skip holiday + week-end
+        self.assertEquals(add(self.datecls(2008, 4, 30), 2), self.datecls(2008, 5, 5))
 
-    def test_add_days_worked(self):
-        self.check_mx()
+    def test_add_days_worked_with_pydatetime(self):
         add = add_days_worked
+        Date = datetime
         # normal
-        self.assertEquals(add(Date(2008, 1, 3), 1), Date(2008, 1, 4))
+        self.assertEquals(add(self.datecls(2008, 1, 3), 1), self.datecls(2008, 1, 4))
         # skip week-end
-        self.assertEquals(add(Date(2008, 1, 3), 2), Date(2008, 1, 7))
+        self.assertEquals(add(self.datecls(2008, 1, 3), 2), self.datecls(2008, 1, 7))
         # skip 2 week-ends
-        self.assertEquals(add(Date(2008, 1, 3), 8), Date(2008, 1, 15))
+        self.assertEquals(add(self.datecls(2008, 1, 3), 8), self.datecls(2008, 1, 15))
         # skip holiday + week-end
-        self.assertEquals(add(Date(2008, 4, 30), 2), Date(2008, 5, 5))
+        self.assertEquals(add(self.datecls(2008, 4, 30), 2), self.datecls(2008, 5, 5))
 
     def test_get_national_holidays(self):
-        self.check_mx()
         holidays = get_national_holidays
-        yield self.assertEquals, holidays(Date(2008, 4, 29), Date(2008, 5, 2)), \
-              [Date(2008, 5, 1)]
-        yield self.assertEquals, holidays(Date(2008, 5, 7), Date(2008, 5, 8)), []
-        x = DateTime(2008, 5, 7, 12, 12, 12)
-        yield self.assertEquals, holidays(x, x+1), []
+        yield self.assertEquals, holidays(self.datecls(2008, 4, 29), self.datecls(2008, 5, 2)), \
+              [self.datecls(2008, 5, 1)]
+        yield self.assertEquals, holidays(self.datecls(2008, 5, 7), self.datecls(2008, 5, 8)), []
+        x = self.datetimecls(2008, 5, 7, 12, 12, 12)
+        yield self.assertEquals, holidays(x, x + self.timedeltacls(days=1)), []
 
     def test_open_days_now_and_before(self):
-        self.check_mx()
         nb = nb_open_days
-        x = now()
-        y = x - RelativeDateTime(seconds=1)
+        x = self.now()
+        y = x - self.timedeltacls(seconds=1)
         self.assertRaises(AssertionError, nb, x, y)
 
-    def check_mx(self):
-        if nb_open_days is None:
-            self.skip('mx.DateTime is not installed')        
-
     def assertOpenDays(self, start, stop, expected):
-        self.check_mx()
         got = nb_open_days(start, stop)
         self.assertEquals(got, expected)
 
     def test_open_days_tuesday_friday(self):
-        self.assertOpenDays(Date(2008, 3, 4), Date(2008, 3, 7), 3)
+        self.assertOpenDays(self.datecls(2008, 3, 4), self.datecls(2008, 3, 7), 3)
 
     def test_open_days_day_nextday(self):
-        self.assertOpenDays(Date(2008, 3, 4), Date(2008, 3, 5), 1)
+        self.assertOpenDays(self.datecls(2008, 3, 4), self.datecls(2008, 3, 5), 1)
 
     def test_open_days_friday_monday(self):
-        self.assertOpenDays(Date(2008, 3, 7), Date(2008, 3, 10), 1)
+        self.assertOpenDays(self.datecls(2008, 3, 7), self.datecls(2008, 3, 10), 1)
 
     def test_open_days_friday_monday_with_two_weekends(self):
-        self.assertOpenDays(Date(2008, 3, 7), Date(2008, 3, 17), 6)
+        self.assertOpenDays(self.datecls(2008, 3, 7), self.datecls(2008, 3, 17), 6)
 
     def test_open_days_tuesday_wednesday(self):
         """week-end + easter monday"""
-        self.assertOpenDays(Date(2008, 3, 18), Date(2008, 3, 26), 5)
+        self.assertOpenDays(self.datecls(2008, 3, 18), self.datecls(2008, 3, 26), 5)
 
     def test_open_days_friday_saturday(self):
-        self.assertOpenDays(Date(2008, 3, 7), Date(2008, 3, 8), 1)
-        
+        self.assertOpenDays(self.datecls(2008, 3, 7), self.datecls(2008, 3, 8), 1)
+
     def test_open_days_friday_sunday(self):
-        self.assertOpenDays(Date(2008, 3, 7), Date(2008, 3, 9), 1)
+        self.assertOpenDays(self.datecls(2008, 3, 7), self.datecls(2008, 3, 9), 1)
 
     def test_open_days_saturday_sunday(self):
-        self.assertOpenDays(Date(2008, 3, 8), Date(2008, 3, 9), 0)
+        self.assertOpenDays(self.datecls(2008, 3, 8), self.datecls(2008, 3, 9), 0)
 
     def test_open_days_saturday_monday(self):
-        self.assertOpenDays(Date(2008, 3, 8), Date(2008, 3, 10), 0)
+        self.assertOpenDays(self.datecls(2008, 3, 8), self.datecls(2008, 3, 10), 0)
 
     def test_open_days_saturday_tuesday(self):
-        self.assertOpenDays(Date(2008, 3, 8), Date(2008, 3, 11), 1)
+        self.assertOpenDays(self.datecls(2008, 3, 8), self.datecls(2008, 3, 11), 1)
 
     def test_open_days_now_now(self):
-        x = now()
+        x = self.now()
         self.assertOpenDays(x, x, 0)
-        
+
     def test_open_days_afternoon_before_holiday(self):
-        self.assertOpenDays(DateTime(2008, 5, 7, 14), Date(2008, 5, 8), 1)
-        
+        self.assertOpenDays(self.datetimecls(2008, 5, 7, 14), self.datetimecls(2008, 5, 8, 0), 1)
+
     def test_open_days_afternoon_before_saturday(self):
-        self.assertOpenDays(DateTime(2008, 5, 9, 14), Date(2008, 5, 10), 1)
-        
+        self.assertOpenDays(self.datetimecls(2008, 5, 9, 14), self.datetimecls(2008, 5, 10, 14), 1)
+
     def test_open_days_afternoon(self):
-        self.assertOpenDays(DateTime(2008, 5, 6, 14), Date(2008, 5, 7), 1)
-        
-    def test_open_days_now_and_one_second(self):
-        x = now()
-        y = x + RelativeDateTime(seconds=1)
-        self.assertOpenDays(x, y, 1)
-        
-    
+        self.assertOpenDays(self.datetimecls(2008, 5, 6, 14), self.datetimecls(2008, 5, 7, 14), 1)
+
+
+class MxDateTC(DateTC):
+    datecls = mxDate
+    datetimecls = mxDateTime
+    timedeltacls = RelativeDateTime
+    now = mxNow
+
+    def check_mx(self):
+        if mxDate is None:
+            self.skip('mx.DateTime is not installed')
+
+    def setUp(self):
+        self.check_mx()
+
+    def test_month(self):
+        """enumerate months"""
+        r = list(date_range(self.datecls(2000,1,2), self.datecls(2000,4,4), endOfMonth))
+        expected = [self.datecls(2000,1,2), self.datecls(2000,2,29), self.datecls(2000,3,31)]
+        self.assertListEquals(r, expected)
+        r = list(date_range(self.datecls(2000,11,30), self.datecls(2001,2,3), endOfMonth))
+        expected = [self.datecls(2000,11,30), self.datecls(2000,12,31), self.datecls(2001,1,31)]
+        self.assertListEquals(r, expected)
+
 if __name__ == '__main__':
     unittest_main()