Repository: peopledoc/workalendar
Pull request #298 (2018-11-29 14:40:59)

Securing the core function arguments to valid types
* only standard library `date` and `datetime` types (or subtypes) are supported,
* added a new dedicated exception,
* added a paragraph in the documentation about this type support.

refs #294


- [x] Tests with a significant number of years to be tested for your calendar.
- [x] Changelog amended with a mention describing your changes.


brunobord: /cc @yassineAlouini and @johnberroa
if you have any input on this PR, please do not hesitate.

Patch:
diff --git a/Changelog.md b/Changelog.md
index 856213b..1d83b8b 100644
--- a/Changelog.md
+++ b/Changelog.md
@@ -3,6 +3,7 @@
 ## master (unreleased)
 
 - Removed dependency to `PyEphem`. This package was the "Python2-compatible" library to deal with the xephem system library. Now it's obsolete, so you don't need this dual-dependency handling, because `ephem` is compatible with Python 2 & Python 3 (#296).
+- Raise an exception when trying to use unsupported date/datetime types (#294).
 
 ## v3.1.1 (2018-11-17)
 
diff --git a/docs/basic.md b/docs/basic.md
index 0c61d38..49a4d53 100644
--- a/docs/basic.md
+++ b/docs/basic.md
@@ -81,10 +81,13 @@ Let's say you want to know how many working days there are between April 2nd and
 50
 ```
 
+## Standard date(time) types only, please!
 
-## Don't worry about the date types!
+For your convenience, we allow both `datetime.date` and `datetime.datetime` types (and their subclasses) when using the core functions.
 
-For your convenience, we allow both `datetime.date` and `datetime.datetime` types when using the core functions. Example:
+**WARNING**: We'll only allow "dates" types coming from the Python standard library. If you're manipulating types from external library. Trying to pass a non-standard argument will result in raising a ``UnsupportedDateType`` error.
+
+Example:
 
 ```python
 >>> from datetime import date, datetime
diff --git a/tox.ini b/tox.ini
index 2121f63..79dce8c 100644
--- a/tox.ini
+++ b/tox.ini
@@ -3,6 +3,7 @@ envlist = py27,flake8,py34,py35,py36,py37,py36-cov,py27-cov
 
 [testenv]
 deps =
+    pandas
     pytest
     cov: pytest-cov
 
diff --git a/workalendar/core.py b/workalendar/core.py
index ab59175..8d8b475 100644
--- a/workalendar/core.py
+++ b/workalendar/core.py
@@ -14,6 +14,8 @@ from calverter import Calverter
 from dateutil import easter
 from lunardate import LunarDate
 
+from .exceptions import UnsupportedDateType
+
 MON, TUE, WED, THU, FRI, SAT, SUN = range(7)
 
 
@@ -27,6 +29,23 @@ class classproperty(object):
         return self.getter(owner)
 
 
+def cleaned_date(day, keep_datetime=False):
+    """
+    Return a "clean" date type.
+
+    * keep a `date` unchanged
+    * convert a datetime into a date,
+    * convert any "duck date" type into a date using its `date()` method.
+    """
+    if not isinstance(day, (date, datetime)):
+        raise UnsupportedDateType(
+            "`{}` is of unsupported type ({})".format(day, type(day)))
+    if not keep_datetime:
+        if hasattr(day, 'date') and callable(day.date):
+            day = day.date()
+    return day
+
+
 class Calendar(object):
 
     FIXED_HOLIDAYS = ()
@@ -79,9 +98,7 @@ class Calendar(object):
 
     def get_holiday_label(self, day):
         """Return the label of the holiday, if the date is a holiday"""
-        # a little exception: chop the datetime type
-        if type(day) is datetime:
-            day = day.date()
+        day = cleaned_date(day)
         return {day: label for day, label in self.holidays(day.year)
                 }.get(day)
 
@@ -117,9 +134,11 @@ class Calendar(object):
         ``extra_holidays`` list.
 
         """
-        # a little exception: chop the datetime type
-        if type(day) is datetime:
-            day = day.date()
+        day = cleaned_date(day)
+        if extra_working_days:
+            extra_working_days = tuple(map(cleaned_date, extra_working_days))
+        if extra_holidays:
+            extra_holidays = tuple(map(cleaned_date, extra_holidays))
 
         # Extra lists exceptions
         if extra_working_days and day in extra_working_days:
@@ -139,9 +158,10 @@ class Calendar(object):
         holidays, even if not in the regular calendar holidays (or weekends).
 
         """
-        # a little exception: chop the datetime type
-        if type(day) is datetime:
-            day = day.date()
+        day = cleaned_date(day)
+
+        if extra_holidays:
+            extra_holidays = tuple(map(cleaned_date, extra_holidays))
 
         if extra_holidays and day in extra_holidays:
             return True
@@ -170,6 +190,14 @@ class Calendar(object):
         Please note that the ``extra_working_days`` list has priority over the
         ``extra_holidays`` list.
         """
+        day = cleaned_date(day, keep_datetime)
+
+        if extra_working_days:
+            extra_working_days = tuple(map(cleaned_date, extra_working_days))
+
+        if extra_holidays:
+            extra_holidays = tuple(map(cleaned_date, extra_holidays))
+
         days = 0
         temp_day = day
         if type(temp_day) is datetime and not keep_datetime:
@@ -221,6 +249,8 @@ class Calendar(object):
         **WARNING**: this function doesn't take into account the calendar
         holidays, only the days of the week and the weekend days parameters.
         """
+        day = cleaned_date(day)
+
         while day.weekday() in self.get_weekend_days():
             day = day + timedelta(days=1)
         return day
@@ -236,6 +266,10 @@ class Calendar(object):
         >>> Calendar.get_nth_weekday_in_month(2013, 1, MON, 2)
         datetime.date(2013, 1, 14)
         """
+        # If start is `None` or Falsy, no need to check and clean
+        if start:
+            start = cleaned_date(start)
+
         day = date(year, month, 1)
         if start:
             day = start
@@ -305,11 +339,8 @@ class Calendar(object):
         This method should even work if your ``start`` and ``end`` arguments
         are datetimes.
         """
-        # Sanitize date types first.
-        if type(start) is datetime:
-            start = start.date()
-        if type(end) is datetime:
-            end = end.date()
+        start = cleaned_date(start)
+        end = cleaned_date(end)
 
         if start == end:
             return 0
diff --git a/workalendar/exceptions.py b/workalendar/exceptions.py
index ac9aa32..7f4bd39 100644
--- a/workalendar/exceptions.py
+++ b/workalendar/exceptions.py
@@ -7,3 +7,9 @@ class CalendarError(Exception):
     """
     Base Calendar Error
     """
+
+
+class UnsupportedDateType(CalendarError):
+    """
+    Raised when trying to use an unsupported date/datetime type.
+    """


Test patch:
diff --git a/workalendar/tests/test_core.py b/workalendar/tests/test_core.py
index 3043c69..6210fc8 100644
--- a/workalendar/tests/test_core.py
+++ b/workalendar/tests/test_core.py
@@ -2,11 +2,14 @@ from datetime import date
 from datetime import datetime
 from unittest import TestCase
 
+import pandas
+
 from workalendar.tests import GenericCalendarTest
 from workalendar.core import MON, TUE, THU, FRI, WED, SAT, SUN
 from workalendar.core import Calendar, LunarCalendar, WesternCalendar
 from workalendar.core import IslamicMixin, JalaliMixin, ChristianMixin
 from workalendar.core import EphemMixin
+from workalendar.exceptions import UnsupportedDateType
 
 
 class CalendarTest(GenericCalendarTest):
@@ -211,78 +214,12 @@ class MockCalendarTest(GenericCalendarTest):
         # test is_holiday
         self.assertTrue(self.cal.is_holiday(christmas))
 
-    def test_add_working_days_datetime(self):
-        # datetime inside, date outside
-        self.assertEquals(
-            self.cal.add_working_days(
-                datetime(self.year, 12, 20, 12, 34, 56), 0),
-            date(self.year, 12, 20)
-        )
-        self.assertEquals(
-            self.cal.add_working_days(
-                datetime(self.year, 12, 20, 12, 34, 56), 1),
-            date(self.year, 12, 21)
-        )
-
-        # Use the `keep_datetime` option
-        self.assertEquals(
-            self.cal.add_working_days(
-                datetime(self.year, 12, 20, 12, 34, 56),
-                0, keep_datetime=True),
-            datetime(self.year, 12, 20, 12, 34, 56)
-        )
-        self.assertEquals(
-            self.cal.add_working_days(
-                datetime(self.year, 12, 20, 12, 34, 56),
-                1, keep_datetime=True),
-            datetime(self.year, 12, 21, 12, 34, 56)
-        )
-
-    def test_sub_working_days_datetime(self):
-        # datetime inside, date outside
-        self.assertEquals(
-            self.cal.sub_working_days(
-                datetime(self.year, 12, 20, 12, 34, 56), 0),
-            date(self.year, 12, 20)
-        )
-        self.assertEquals(
-            self.cal.sub_working_days(
-                datetime(self.year, 12, 20, 12, 34, 56), 1),
-            date(self.year, 12, 19)
-        )
-
-        # Use the `keep_datetime` option
-        self.assertEquals(
-            self.cal.sub_working_days(
-                datetime(self.year, 12, 20, 12, 34, 56),
-                0, keep_datetime=True),
-            datetime(self.year, 12, 20, 12, 34, 56)
-        )
-        self.assertEquals(
-            self.cal.sub_working_days(
-                datetime(self.year, 12, 20, 12, 34, 56),
-                1, keep_datetime=True),
-            datetime(self.year, 12, 19, 12, 34, 56)
-        )
-
-    def test_datetime(self):
-        self.assertFalse(
-            self.cal.is_working_day(datetime(2014, 1, 1)))
-        self.assertTrue(
-            self.cal.is_holiday(datetime(2014, 1, 1)))
-
     def test_get_holiday_label(self):
         self.assertEqual(
             self.cal.get_holiday_label(date(2014, 1, 1)), 'New year')
         self.assertIsNone(
             self.cal.get_holiday_label(date(2014, 1, 2)))
 
-    def test_get_holiday_label_with_datetime(self):
-        self.assertEqual(
-            self.cal.get_holiday_label(datetime(2014, 1, 1)), 'New year')
-        self.assertIsNone(
-            self.cal.get_holiday_label(datetime(2014, 1, 2)))
-
     def test_add_working_days_backwards(self):
         day = date(self.year, 1, 3)
         # since this calendar has no weekends, we'll just have a 1-day-shift
@@ -577,3 +514,242 @@ class CalendarClassName(TestCase):
         self.assertEqual(
             MultipleLineEmptyFirstDocstring.name, "Multiple line empty first"
         )
+
+
+class TestAcceptableDateTypes(GenericCalendarTest):
+    """
+    Test cases about accepted date and datetime types.
+    """
+    cal_class = MockCalendar
+    unsupported = ('hello', 1)
+
+    def test_unsupported_type_is_working_day(self):
+        for arg in self.unsupported:
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.is_working_day(arg)
+
+        # Extra holidays optional argument
+        for arg in self.unsupported:
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.is_working_day(
+                    date(2018, 1, 1),
+                    extra_holidays=[arg]
+                )
+        # Extra working days optional argument
+        for arg in self.unsupported:
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.is_working_day(
+                    date(2018, 1, 1),
+                    extra_working_days=[arg]
+                )
+
+    def test_unsupported_type_is_holiday(self):
+        for arg in self.unsupported:
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.is_holiday(arg)
+
+        # Extra holidays optional argument
+        for arg in self.unsupported:
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.is_holiday(
+                    date(2018, 1, 1),
+                    extra_holidays=[arg]
+                )
+
+    def test_unsupported_type_holiday_label(self):
+        for arg in self.unsupported:
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.get_holiday_label(arg)
+
+    def test_unsupported_type_add_sub_working_days(self):
+        for arg in self.unsupported:
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.add_working_days(arg, 1)
+
+        for arg in self.unsupported:
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.sub_working_days(arg, 1)
+
+        # Extra holidays optional argument
+        for arg in self.unsupported:
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.add_working_days(
+                    date(2018, 1, 1), 1,
+                    extra_holidays=[arg]
+                )
+        # Extra working days optional argument
+        for arg in self.unsupported:
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.add_working_days(
+                    date(2018, 1, 1), 1,
+                    extra_working_days=[arg]
+                )
+        # NOTE: no need to test "sub", they're calling each other.
+
+    def test_unsupported_type_find_following_working_day(self):
+        for arg in self.unsupported:
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.find_following_working_day(arg)
+
+    def test_unsupported_type_get_nth_weekday_in_month(self):
+        for arg in self.unsupported:
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.get_nth_weekday_in_month(2018, 1, MON, start=arg)
+
+    def test_unsupported_type_get_working_days_delta(self):
+        for arg in self.unsupported:
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.get_working_days_delta(date(2018, 1, 1), arg)
+
+            with self.assertRaises(UnsupportedDateType):
+                self.cal.get_working_days_delta(arg, date(2018, 1, 1))
+
+    def test_datetime(self):
+        self.assertFalse(
+            self.cal.is_working_day(datetime(2014, 1, 1)))
+        self.assertTrue(
+            self.cal.is_holiday(datetime(2014, 1, 1)))
+
+    def test_add_working_days_datetime(self):
+        # datetime inside, date outside
+        self.assertEquals(
+            self.cal.add_working_days(
+                datetime(self.year, 12, 20, 12, 34, 56), 0),
+            date(self.year, 12, 20)
+        )
+        self.assertEquals(
+            self.cal.add_working_days(
+                datetime(self.year, 12, 20, 12, 34, 56), 1),
+            date(self.year, 12, 21)
+        )
+
+        # Use the `keep_datetime` option
+        self.assertEquals(
+            self.cal.add_working_days(
+                datetime(self.year, 12, 20, 12, 34, 56),
+                0, keep_datetime=True),
+            datetime(self.year, 12, 20, 12, 34, 56)
+        )
+        self.assertEquals(
+            self.cal.add_working_days(
+                datetime(self.year, 12, 20, 12, 34, 56),
+                1, keep_datetime=True),
+            datetime(self.year, 12, 21, 12, 34, 56)
+        )
+
+    def test_sub_working_days_datetime(self):
+        # datetime inside, date outside
+        self.assertEquals(
+            self.cal.sub_working_days(
+                datetime(self.year, 12, 20, 12, 34, 56), 0),
+            date(self.year, 12, 20)
+        )
+        self.assertEquals(
+            self.cal.sub_working_days(
+                datetime(self.year, 12, 20, 12, 34, 56), 1),
+            date(self.year, 12, 19)
+        )
+
+        # Use the `keep_datetime` option
+        self.assertEquals(
+            self.cal.sub_working_days(
+                datetime(self.year, 12, 20, 12, 34, 56),
+                0, keep_datetime=True),
+            datetime(self.year, 12, 20, 12, 34, 56)
+        )
+        self.assertEquals(
+            self.cal.sub_working_days(
+                datetime(self.year, 12, 20, 12, 34, 56),
+                1, keep_datetime=True),
+            datetime(self.year, 12, 19, 12, 34, 56)
+        )
+
+    def test_get_holiday_label_with_datetime(self):
+        self.assertEqual(
+            self.cal.get_holiday_label(datetime(2014, 1, 1)), 'New year')
+        self.assertIsNone(
+            self.cal.get_holiday_label(datetime(2014, 1, 2)))
+
+
+class PandasTimestampTest(GenericCalendarTest):
+    cal_class = MockCalendar
+
+    def test_panda_type_is_working_day(self):
+        self.assertFalse(
+            self.cal.is_working_day(pandas.to_datetime("2018-1-1"))
+        )
+
+        # Extra holidays optional argument
+        self.assertFalse(
+            self.cal.is_working_day(
+                date(2018, 1, 2),
+                extra_holidays=[pandas.to_datetime("2018-1-2")]
+            )
+        )
+        # Extra working days optional argument
+        self.assertTrue(
+            self.cal.is_working_day(
+                date(2018, 1, 1),
+                extra_working_days=[pandas.to_datetime("2018-1-1")]
+            )
+        )
+
+    def test_panda_type_is_holiday(self):
+        self.assertTrue(self.cal.is_holiday(pandas.to_datetime("2018-1-1")))
+
+        # Extra holidays optional argument
+        self.assertTrue(
+            self.cal.is_holiday(
+                date(2018, 2, 1),
+                extra_holidays=[pandas.to_datetime("2018-2-1")]
+            )
+        )
+
+    def test_panda_type_holiday_label(self):
+        label = self.cal.get_holiday_label(pandas.to_datetime("2018-1-1"))
+        self.assertEqual(label, "New year")
+
+    def test_panda_type_add_sub_working_days(self):
+        day = pandas.to_datetime("2018-12-24")
+        next_day = self.cal.add_working_days(day, 1)
+        self.assertEqual(next_day, date(2018, 12, 26))
+
+        previous_day = self.cal.sub_working_days(next_day, 1)
+        self.assertEqual(previous_day, date(2018, 12, 24))
+
+        next_day = self.cal.add_working_days(
+            date(2018, 12, 24), 1,
+            extra_holidays=[pandas.to_datetime("2018-12-26")]
+        )
+        self.assertEqual(next_day, date(2018, 12, 27))
+
+        next_day = self.cal.add_working_days(
+            date(2018, 12, 24), 1,
+            extra_working_days=[pandas.to_datetime("2018-12-25")]
+        )
+        self.assertEqual(next_day, date(2018, 12, 25))
+
+    def test_unsupported_type_find_following_working_day(self):
+        following_day = self.cal.find_following_working_day(
+            pandas.to_datetime("2018-1-1")
+        )
+        # No weekend days, the next day is "today"
+        self.assertEqual(following_day, date(2018, 1, 1))
+
+    def test_unsupported_type_get_nth_weekday_in_month(self):
+        start = pandas.to_datetime("2018-1-4")
+        monday = self.cal.get_nth_weekday_in_month(2018, 1, MON, start=start)
+        self.assertEqual(monday, date(2018, 1, 8))
+
+    def test_unsupported_type_get_working_days_delta(self):
+        start, end = date(2018, 12, 23), pandas.to_datetime("2018-12-26")
+        delta = self.cal.get_working_days_delta(start, end)
+        self.assertEqual(delta, 2)
+        delta = self.cal.get_working_days_delta(end, start)
+        self.assertEqual(delta, 2)
+
+        start, end = pandas.to_datetime("2018-12-23"), date(2018, 12, 26)
+        delta = self.cal.get_working_days_delta(start, end)
+        self.assertEqual(delta, 2)
+        delta = self.cal.get_working_days_delta(end, start)
+        self.assertEqual(delta, 2)