Merge pull request #31388 from dimagi/sk/date_add

case search "date_add" function

corehq/apps/case_search/dsl_utils.py

@@ -9,22 +9,22 @@
 from corehq.apps.case_search.xpath_functions import XPATH_VALUE_FUNCTIONS
 
 
-def unwrap_value(node, context):
+def unwrap_value(value, context):
     """Returns the value of the node if it is wrapped in a function, otherwise just returns the node
     """
-    if isinstance(node, UnaryExpression) and node.op == '-':
-        return -1 * node.right
-    if not isinstance(node, FunctionCall):
-        return node
+    if isinstance(value, UnaryExpression) and value.op == '-':
+        return -1 * value.right
+    if not isinstance(value, FunctionCall):
+        return value
     try:
-        return XPATH_VALUE_FUNCTIONS[node.name](node, context)
+        return XPATH_VALUE_FUNCTIONS[value.name](value, context)
     except KeyError:
         raise CaseFilterError(
             _("We don't know what to do with the function \"{}\". Accepted functions are: {}").format(
-                node.name,
+                value.name,
                 ", ".join(list(XPATH_VALUE_FUNCTIONS.keys())),
             ),
-            serialize(node)
+            serialize(value)
         )
     except XPathFunctionException as e:
-        raise CaseFilterError(str(e), serialize(node))
+        raise CaseFilterError(str(e), serialize(value))

corehq/apps/case_search/tests/test_value_functions.py

@@ -2,11 +2,15 @@
 
 from eulxml.xpath import parse as parse_xpath
 from freezegun import freeze_time
-from testil import eq
+from testil import assert_raises, eq
 
 from corehq.apps.case_search.exceptions import XPathFunctionException
 from corehq.apps.case_search.filter_dsl import SearchFilterContext
-from corehq.apps.case_search.xpath_functions.value_functions import today, date
+from corehq.apps.case_search.xpath_functions.value_functions import (
+    date,
+    date_add,
+    today,
+)
 from corehq.apps.domain.models import Domain
 from corehq.apps.domain.shortcuts import create_domain
 
@@ -57,3 +61,64 @@ def test_date_today(self):
         node = parse_xpath("date(today())")
         result = date(node, SearchFilterContext("domain"))
         eq(result, '2021-08-02')
+
+
+def test_date_add():
+    def _do_test(expression, expected):
+        node = parse_xpath(expression)
+        result = date_add(node, SearchFilterContext("domain"))
+        eq(result, expected)
+
+    test_cases = [
+        ("date_add('2020-02-29', 'years', 1)", "2021-02-28"),
+        ("date_add('2020-02-29', 'years', 1.0)", "2021-02-28"),
+        ("date_add('2022-01-01', 'months', 1)", "2022-02-01"),
+        ("date_add('2022-01-01', 'months', '3')", "2022-04-01"),
+        ("date_add('2020-04-30', 'months', -2)", "2020-02-29"),
+        ("date_add('2020-03-31', 'weeks', 4)", "2020-04-28"),
+        ("date_add('2020-02-01', 'weeks', 1.5)", "2020-02-11"),
+        ("date_add('2022-01-01', 'days', -1)", "2021-12-31"),
+        ("date_add('2022-01-01', 'days', 1.5)", "2022-01-02"),
+        ("date_add('2022-01-01', 'days', '5')", "2022-01-06"),
+        ("date_add(0, 'days', '5')", "1970-01-06"),
+        ("date_add(365, 'years', '5')", "1976-01-01"),
+        ("date_add(date('2021-01-01'), 'years', '5')", "2026-01-01"),
+        ("date_add(date(5), 'months', '1')", "1970-02-06"),
+    ]
+
+    for expression, expected in test_cases:
+        yield _do_test, expression, expected
+
+
+def test_date_add_errors():
+    def _do_test(expression):
+        node = parse_xpath(expression)
+        with assert_raises(XPathFunctionException):
+            date_add(node, SearchFilterContext("domain"))
+
+    test_cases = [
+        # bad date
+        "date_add('2020-02-31', 'years', 1)",
+        "date_add(1.5, 'years', 1)",
+
+        # non-integer quantity
+        "date_add('2020-02-01', 'years', 1.5)",
+        "date_add('2020-02-01', 'months', 1.5)",
+
+        # non-numeric quantity
+        "date_add('2020-02-01', 'months', 'five')",
+
+        # bad interval names
+        "date_add('2020-02-01', 'year', 1)",
+        "date_add('2020-02-01', 'month', 1)",
+        "date_add('2020-02-01', 'week', 1)",
+        "date_add('2020-02-01', 'day', 1)",
+
+        # missing args
+        "date_add('2020-02-01', 'days')",
+        "date_add('2020-02-01')",
+        "date_add()",
+    ]
+
+    for expression in test_cases:
+        yield _do_test, expression

corehq/apps/case_search/xpath_functions/__init__.py

@@ -1,10 +1,11 @@
 from .query_functions import not_, selected_all, selected_any
 from .subcase_functions import subcase
-from .value_functions import date, today
+from .value_functions import date, date_add, today
 
 # functions that transform or produce a value
 XPATH_VALUE_FUNCTIONS = {
     'date': date,
+    'date_add': date_add,
     'today': today,
 }
 

corehq/apps/case_search/xpath_functions/value_functions.py

@@ -1,14 +1,16 @@
 import datetime
 
-import pytz
 from django.utils.dateparse import parse_date
 from django.utils.translation import gettext as _
 
+import pytz
+from dateutil.relativedelta import relativedelta
 from eulxml.xpath.ast import serialize
 
+from dimagi.utils.parsing import ISO_DATE_FORMAT
+
 from corehq.apps.case_search.exceptions import XPathFunctionException
 from corehq.apps.domain.models import Domain
-from dimagi.utils.parsing import ISO_DATE_FORMAT
 
 
 def date(node, context):
@@ -24,27 +26,30 @@ def date(node, context):
 
     arg = unwrap_value(arg, context)
 
-    if isinstance(arg, int):
-        return (datetime.date(1970, 1, 1) + datetime.timedelta(days=arg)).strftime("%Y-%m-%d")
+    parsed_date = _value_to_date(node, arg)
+    return parsed_date.strftime(ISO_DATE_FORMAT)
 
-    if isinstance(arg, str):
+
+def _value_to_date(node, value):
+    if isinstance(value, int):
+        parsed_date = datetime.date(1970, 1, 1) + datetime.timedelta(days=value)
+    elif isinstance(value, str):
         try:
-            parsed_date = parse_date(arg)
+            parsed_date = parse_date(value)
         except ValueError:
-            raise XPathFunctionException(_("{} is not a valid date").format(arg), serialize(node))
-
-        if parsed_date is None:
-            raise XPathFunctionException(
-                _("The \"date\" function only accepts strings of the format \"YYYY-mm-dd\""),
-                serialize(node)
-            )
+            raise XPathFunctionException(_("{} is not a valid date").format(value), serialize(node))
+    elif isinstance(value, datetime.date):
+        parsed_date = value
+    else:
+        parsed_date = None
 
-        return arg
+    if parsed_date is None:
+        raise XPathFunctionException(
+            _("Invalid date value. Dates must be an integer or a string of the format \"YYYY-mm-dd\""),
+            serialize(node)
+        )
 
-    raise XPathFunctionException(
-        "The \"date\" function only accepts integers or strings of the format \"YYYY-mm-dd\"",
-        serialize(node)
-    )
+    return parsed_date
 
 
 def today(node, context):
@@ -58,3 +63,57 @@ def today(node, context):
     domain_obj = Domain.get_by_name(context.domain)
     timezone = domain_obj.get_default_timezone() if domain_obj else pytz.UTC
     return datetime.datetime.now(timezone).strftime(ISO_DATE_FORMAT)
+
+
+def date_add(node, context):
+    from corehq.apps.case_search.dsl_utils import unwrap_value
+
+    assert node.name == 'date_add'
+    if len(node.args) != 3:
+        raise XPathFunctionException(
+            _("The \"date_add\" function expects three arguments, got {count}").format(count=len(node.args)),
+            serialize(node)
+        )
+
+    date_arg = unwrap_value(node.args[0], context)
+    date_value = _value_to_date(node, date_arg)
+
+    interval_type = unwrap_value(node.args[1], context)
+    interval_types = ("days", "weeks", "months", "years")
+    if interval_type not in interval_types:
+        raise XPathFunctionException(
+            _("The \"date_add\" function expects the 'interval' argument to be one of {types}").format(
+                types=interval_types
+            ),
+            serialize(node)
+        )
+
+    quantity = unwrap_value(node.args[2], context)
+    if isinstance(quantity, str):
+        try:
+            quantity = float(quantity)
+        except (ValueError, TypeError):
+            raise XPathFunctionException(
+                _("The \"date_add\" function expects the interval quantity to be a numeric value"),
+                serialize(node)
+            )
+
+    if not isinstance(quantity, (int, float)):
+        raise XPathFunctionException(
+            _("The \"date_add\" function expects the interval quantity to be a numeric value"),
+            serialize(node)
+        )
+
+    if interval_type in ("years", "months") and int(quantity) != quantity:
+        raise XPathFunctionException(
+            _("Non-integer years and months are ambiguous and not supported by the \"date_add\" function"),
+            serialize(node)
+        )
+
+    try:
+        result = date_value + relativedelta(**{interval_type: quantity})
+    except Exception as e:
+        # catchall in case of an unexpected error
+        raise XPathFunctionException(str(e), serialize(node))
+
+    return result.strftime(ISO_DATE_FORMAT)