dateutil_dates_between() function, refs #1

simonw · Sep 29, 2020 · adf1b65 · adf1b65
1 parent 1f96183
commit adf1b65
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 0 deletions.
diff --git a/datasette_dateutil/__init__.py b/datasette_dateutil/__init__.py
@@ -2,6 +2,7 @@
 from dateutil.parser import parse, ParserError
 from dateutil.rrule import rrulestr
 from dateutil.easter import easter
+import datetime
 import itertools
 import json
 
@@ -69,6 +70,30 @@ def dateutil_rrule_date(rrule, dtstart=None):
     return dateutil_rrule(rrule, dtstart, date=True)
 
 
+def _between(dt_start, dt_end, inclusive=True):
+    if dt_start >= dt_end:
+        return []
+    current = dt_start
+    while current < dt_end:
+        yield current.isoformat()
+        current += datetime.timedelta(days=1)
+    if inclusive:
+        yield current.isoformat()
+
+
+def dateutil_dates_between(start, end, inclusive=True):
+    dt_start = parse(start).date()
+    dt_end = parse(end).date()
+    results = list(
+        itertools.islice(_between(dt_start, dt_end, inclusive), 0, RRULE_MAX + 1)
+    )
+    if len(results) > RRULE_MAX:
+        raise TooManyError(
+            "More than {} dates between '{}' and '{}".format(RRULE_MAX, start, end)
+        )
+    return json.dumps(results)
+
+
 @hookimpl
 def prepare_connection(conn):
     conn.create_function("dateutil_parse", 1, dateutil_parse)
@@ -82,3 +107,5 @@ def prepare_connection(conn):
     conn.create_function("dateutil_rrule", 2, dateutil_rrule)
     conn.create_function("dateutil_rrule_date", 1, dateutil_rrule_date)
     conn.create_function("dateutil_rrule_date", 2, dateutil_rrule_date)
+    conn.create_function("dateutil_dates_between", 2, dateutil_dates_between)
+    conn.create_function("dateutil_dates_between", 3, dateutil_dates_between)
diff --git a/tests/test_dateutil.py b/tests/test_dateutil.py
@@ -50,6 +50,14 @@ async def test_plugin_is_installed():
             "select dateutil_rrule_date('FREQ=DAILY;INTERVAL=10;COUNT=5', '2020-01-01')",
             '["2020-01-01", "2020-01-11", "2020-01-21", "2020-01-31", "2020-02-10"]',
         ),
+        (
+            "select dateutil_dates_between('1 january 2020', '5 jan 2020', 0)",
+            '["2020-01-01", "2020-01-02", "2020-01-03", "2020-01-04"]',
+        ),
+        (
+            "select dateutil_dates_between('1 january 2020', '5 jan 2020', 1)",
+            '["2020-01-01", "2020-01-02", "2020-01-03", "2020-01-04", "2020-01-05"]',
+        ),
     ],
 )
 async def test_dateutil_sql_functions(sql, expected):