Add a function to compute short statutory time periods (requiring at …

…least n business days) (#124) * Fix #122, function to return first date after at least n business days * Added a few more tests Co-authored-by: nonprofittechy <[email protected]> Co-authored-by: Bryce Willey <[email protected]>
SuffolkLITLab · Nov 3, 2022 · e3fa243 · e3fa243
1 parent 9972330
commit e3fa243
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 2 deletions.
diff --git a/docassemble/ALToolbox/business_days.py b/docassemble/ALToolbox/business_days.py
@@ -12,6 +12,14 @@
   3. https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html  
 """
 
+__all__ = [
+    "standard_holidays",
+    "non_business_days",
+    "is_business_day",
+    "get_next_business_day",
+    "get_date_after_n_business_days",
+]
+
 
 def standard_holidays(
     year,
@@ -25,6 +33,8 @@ def standard_holidays(
     Note that this draws on the "holidays" package which may deviate slightly from
     holidays observed by a local court, but should be very close to accurate.
 
+    add_holidays should be a dictionary from dates ("12-15") to the name of the holiday.
+
     Returns a dictionary like-object that you can treat like:
     {
         "2021-01-01": "New Year's Day",
@@ -201,3 +211,31 @@ def get_next_business_day(
     ):
         date_to_check = date_to_check.plus(days=1)
     return date_to_check
+
+
+def get_date_after_n_business_days(
+    start_date: Union[str, DADateTime],
+    wait_n_days=1,
+    country="US",
+    subdiv="MA",
+    add_holidays: Mapping = None,
+    remove_holidays: Iterable[str] = None,
+) -> DADateTime:
+    """
+    Returns a time period which contains a minimum of `n` business days.
+    """
+    if not isinstance(start_date, DADateTime):
+        start_date = as_datetime(start_date)
+    date_to_check = start_date
+
+    for _ in range(wait_n_days):
+        date_to_check = date_to_check.plus(days=1)
+        while not is_business_day(
+            date_to_check,
+            country=country,
+            subdiv=subdiv,
+            add_holidays=add_holidays,
+            remove_holidays=remove_holidays,
+        ):
+            date_to_check = date_to_check.plus(days=1)
+    return date_to_check
diff --git a/docassemble/ALToolbox/test_altoolbox.py b/docassemble/ALToolbox/test_altoolbox.py
@@ -1,6 +1,10 @@
 import unittest
-from .business_days import is_business_day, get_next_business_day
-from docassemble.base.util import today
+from .business_days import (
+    is_business_day,
+    get_next_business_day,
+    get_date_after_n_business_days,
+)
+from docassemble.base.util import today, as_datetime
 
 
 class TestBusinessDays(unittest.TestCase):
@@ -10,6 +14,28 @@ def test_is_business_day(self):
         this_years_christmas = today().replace(month=12, day=25)
         self.assertFalse(is_business_day(this_years_christmas))
 
+    def test_get_date_after_n_business_days(self):
+        self.assertEqual(
+            get_date_after_n_business_days("2022-12-24", 5), as_datetime("2023-01-03")
+        )
+        self.assertEqual(
+            get_date_after_n_business_days("2022-12-12", 3), as_datetime("2022-12-15")
+        )
+        self.assertEqual(
+            get_date_after_n_business_days(
+                "2022-12-12",
+                3,
+                add_holidays={"12-13": "fake day 1", "12-14": "fake day2"},
+            ),
+            as_datetime("2022-12-19"),
+        )
+        self.assertEqual(
+            get_date_after_n_business_days(
+                "2022-12-24", 5, remove_holidays=["Christmas Day"]
+            ),
+            as_datetime("2022-12-30"),
+        )
+
 
 if __name__ == "__main__":
     unittest.main()