Optional dtstart second arg for dateutil_rrule() and dateutil_rrule_d…

…ate(), refs #1
simonw · Sep 29, 2020 · 1f96183 · 1f96183
1 parent 686c683
commit 1f96183
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 4 deletions.
diff --git a/datasette_dateutil/__init__.py b/datasette_dateutil/__init__.py
@@ -50,8 +50,12 @@ class TooManyError(Exception):
     pass
 
 
-def dateutil_rrule(rrule, date=False):
-    results = list(itertools.islice(rrulestr(rrule), 0, RRULE_MAX + 1))
+def dateutil_rrule(rrule, dtstart=None, date=False):
+    dtstart = dtstart and parse(dtstart) or None
+    kwargs = {}
+    if dtstart:
+        kwargs["dtstart"] = dtstart
+    results = list(itertools.islice(rrulestr(rrule, **kwargs), 0, RRULE_MAX + 1))
     if len(results) > RRULE_MAX:
         raise TooManyError(
             "More than {} results returned by '{}'".format(RRULE_MAX, rrule)
@@ -61,8 +65,8 @@ def dateutil_rrule(rrule, date=False):
     return json.dumps([d.isoformat() for d in results])
 
 
-def dateutil_rrule_date(rrule):
-    return dateutil_rrule(rrule, date=True)
+def dateutil_rrule_date(rrule, dtstart=None):
+    return dateutil_rrule(rrule, dtstart, date=True)
 
 
 @hookimpl
@@ -75,4 +79,6 @@ def prepare_connection(conn):
     )
     conn.create_function("dateutil_easter", 1, dateutil_easter)
     conn.create_function("dateutil_rrule", 1, dateutil_rrule)
+    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)
diff --git a/tests/test_dateutil.py b/tests/test_dateutil.py
@@ -38,10 +38,18 @@ async def test_plugin_is_installed():
             "select dateutil_rrule('DTSTART:20200101\nFREQ=DAILY;INTERVAL=10;COUNT=5')",
             '["2020-01-01T00:00:00", "2020-01-11T00:00:00", "2020-01-21T00:00:00", "2020-01-31T00:00:00", "2020-02-10T00:00:00"]',
         ),
+        (
+            "select dateutil_rrule('FREQ=DAILY;INTERVAL=10;COUNT=5', '2020-01-01')",
+            '["2020-01-01T00:00:00", "2020-01-11T00:00:00", "2020-01-21T00:00:00", "2020-01-31T00:00:00", "2020-02-10T00:00:00"]',
+        ),
         (
             "select dateutil_rrule_date('DTSTART:20200101\nFREQ=DAILY;INTERVAL=10;COUNT=5')",
             '["2020-01-01", "2020-01-11", "2020-01-21", "2020-01-31", "2020-02-10"]',
         ),
+        (
+            "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"]',
+        ),
     ],
 )
 async def test_dateutil_sql_functions(sql, expected):