dateutil_rrule and dateutil_rrule_date functions, refs #1

simonw · Sep 29, 2020 · 686c683 · 686c683
1 parent 0095bc8
commit 686c683
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 17 deletions.
diff --git a/datasette_dateutil/__init__.py b/datasette_dateutil/__init__.py
@@ -1,6 +1,12 @@
 from datasette import hookimpl
 from dateutil.parser import parse, ParserError
+from dateutil.rrule import rrulestr
 from dateutil.easter import easter
+import itertools
+import json
+
+
+RRULE_MAX = 10_000
 
 
 def _dateutil_parse_shared(s, **kwargs):
@@ -40,6 +46,25 @@ def dateutil_easter(year):
         return None
 
 
+class TooManyError(Exception):
+    pass
+
+
+def dateutil_rrule(rrule, date=False):
+    results = list(itertools.islice(rrulestr(rrule), 0, RRULE_MAX + 1))
+    if len(results) > RRULE_MAX:
+        raise TooManyError(
+            "More than {} results returned by '{}'".format(RRULE_MAX, rrule)
+        )
+    if date:
+        results = [d.date() for d in results]
+    return json.dumps([d.isoformat() for d in results])
+
+
+def dateutil_rrule_date(rrule):
+    return dateutil_rrule(rrule, date=True)
+
+
 @hookimpl
 def prepare_connection(conn):
     conn.create_function("dateutil_parse", 1, dateutil_parse)
@@ -49,3 +74,5 @@ def prepare_connection(conn):
         "dateutil_parse_fuzzy_dayfirst", 1, dateutil_parse_fuzzy_dayfirst
     )
     conn.create_function("dateutil_easter", 1, dateutil_easter)
+    conn.create_function("dateutil_rrule", 1, dateutil_rrule)
+    conn.create_function("dateutil_rrule_date", 1, dateutil_rrule_date)
diff --git a/tests/test_dateutil.py b/tests/test_dateutil.py
@@ -14,26 +14,60 @@ async def test_plugin_is_installed():
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("sql,expected", [
-    ("select dateutil_parse('1st october 2009')", "2009-10-01T00:00:00"),
-    ("select dateutil_parse('invalid')", None),
-    ("select dateutil_parse('due on 1st october 2009')", None),
-    ("select dateutil_parse_fuzzy('due on 1st october 2009')", "2009-10-01T00:00:00"),
-    ("select dateutil_parse_fuzzy('due on')", None),
-    ("select dateutil_parse_dayfirst('1/2/2020')", "2020-02-01T00:00:00"),
-    ("select dateutil_parse('1/2/2020')", "2020-01-02T00:00:00"),
-    ("select dateutil_parse_fuzzy('due on 1/2/2003')", "2003-01-02T00:00:00"),
-    ("select dateutil_parse_fuzzy_dayfirst('due on 1/2/2003')", "2003-02-01T00:00:00"),
-    ("select dateutil_easter(2020)", "2020-04-12"),
-    ("select dateutil_easter('invalid')", None),
-])
+@pytest.mark.parametrize(
+    "sql,expected",
+    [
+        ("select dateutil_parse('1st october 2009')", "2009-10-01T00:00:00"),
+        ("select dateutil_parse('invalid')", None),
+        ("select dateutil_parse('due on 1st october 2009')", None),
+        (
+            "select dateutil_parse_fuzzy('due on 1st october 2009')",
+            "2009-10-01T00:00:00",
+        ),
+        ("select dateutil_parse_fuzzy('due on')", None),
+        ("select dateutil_parse_dayfirst('1/2/2020')", "2020-02-01T00:00:00"),
+        ("select dateutil_parse('1/2/2020')", "2020-01-02T00:00:00"),
+        ("select dateutil_parse_fuzzy('due on 1/2/2003')", "2003-01-02T00:00:00"),
+        (
+            "select dateutil_parse_fuzzy_dayfirst('due on 1/2/2003')",
+            "2003-02-01T00:00:00",
+        ),
+        ("select dateutil_easter(2020)", "2020-04-12"),
+        ("select dateutil_easter('invalid')", None),
+        (
+            "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_date('DTSTART:20200101\nFREQ=DAILY;INTERVAL=10;COUNT=5')",
+            '["2020-01-01", "2020-01-11", "2020-01-21", "2020-01-31", "2020-02-10"]',
+        ),
+    ],
+)
 async def test_dateutil_sql_functions(sql, expected):
     app = Datasette([], memory=True).app()
     async with httpx.AsyncClient(app=app) as client:
-        response = await client.get("http://localhost/:memory:.json", params={
-            "sql": sql,
-            "_shape": "array",
-        })
+        response = await client.get(
+            "http://localhost/:memory:.json",
+            params={
+                "sql": sql,
+                "_shape": "array",
+            },
+        )
         assert 200 == response.status_code
         actual = list(response.json()[0].values())[0]
         assert actual == expected
+
+
+@pytest.mark.asyncio
+async def test_dateutil_unbounded_rrule_error():
+    app = Datasette([], memory=True).app()
+    async with httpx.AsyncClient(app=app) as client:
+        response = await client.get(
+            "http://localhost/:memory:.json",
+            params={
+                "sql": "select dateutil_rrule('FREQ=DAILY;INTERVAL=10')",
+                "_shape": "array",
+            },
+        )
+        assert 400 == response.status_code