Repository: peopledoc/workalendar
Pull request #270 (2018-07-05 14:18:55)

Add working days using datetimes

- [x] Tests with a significant number of years to be tested for your calendar.
- [x] Changelog amended with a mention describing your changes.



Patch:
diff --git a/CHANGELOG b/CHANGELOG
index f46e913..c4920e0 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -4,7 +4,8 @@ CHANGELOG
 master (unreleased)
 -------------------
 
-Nothing here yet.
+- Allow the `add_working_days()` function to be provided a datetime, and returning a `date` (#270).
+- Added a `keep_datetime` option to keep the original type of the input argument for both ``add_working_days()`` and ``sub_working_days()`` functions (#270).
 
 2.5.0 (2018-06-14)
 ------------------
diff --git a/workalendar/core.py b/workalendar/core.py
index 9121d31..41a3ad7 100644
--- a/workalendar/core.py
+++ b/workalendar/core.py
@@ -130,9 +130,14 @@ class Calendar(object):
         return day in self.holidays_set(day.year)
 
     def add_working_days(self, day, delta,
-                         extra_working_days=None, extra_holidays=None):
+                         extra_working_days=None, extra_holidays=None,
+                         keep_datetime=False):
         """Add `delta` working days to the date.
 
+        You can provide either a date or a datetime to this function that will
+        output a ``date`` result. You can alter this behaviour using the
+        ``keep_datetime`` option set to ``True``.
+
         the ``delta`` parameter might be positive or negative. If it's
         negative, you may want to use the ``sub_working_days()`` method with
         a positive ``delta`` argument.
@@ -148,6 +153,8 @@ class Calendar(object):
         """
         days = 0
         temp_day = day
+        if type(temp_day) is datetime and not keep_datetime:
+            temp_day = temp_day.date()
         day_added = 1 if delta >= 0 else -1
         delta = abs(delta)
         while days < delta:
@@ -159,7 +166,8 @@ class Calendar(object):
         return temp_day
 
     def sub_working_days(self, day, delta,
-                         extra_working_days=None, extra_holidays=None):
+                         extra_working_days=None, extra_holidays=None,
+                         keep_datetime=False):
         """
         Substract `delta` working days to the date.
 
@@ -177,10 +185,16 @@ class Calendar(object):
             cal.sub_working_days(my_date, -7)
             cal.sub_working_days(my_date, 7)
 
+        As in ``add_working_days()`` you can set the parameter
+        ``keep_datetime`` to ``True`` to make sure that if your ``day``
+        argument is a ``datetime``, the returned date will also be a
+        ``datetime`` object.
+
         """
         delta = abs(delta)
         return self.add_working_days(
-            day, -delta, extra_working_days, extra_holidays)
+            day, -delta,
+            extra_working_days, extra_holidays, keep_datetime=keep_datetime)
 
     def find_following_working_day(self, day):
         "Looks for the following working day"


Test patch:
diff --git a/workalendar/tests/test_core.py b/workalendar/tests/test_core.py
index 57aa7f8..a869ba4 100644
--- a/workalendar/tests/test_core.py
+++ b/workalendar/tests/test_core.py
@@ -211,6 +211,60 @@ 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)))