Merge pull request #229 from opensafely/support-arithmetic-operators

Allow arithmetic operators in SQL expressions

cohortextractor/VERSION

@@ -1 +1 @@
-0.1.28
+0.1.29

cohortextractor/expressions.py

@@ -112,6 +112,8 @@ def is_allowed(token):
         return value in [">", "<", ">=", "<=", "=", "!="]
     if ttype in ttypes.Number.Float or ttype in ttypes.Number.Integer:
         return True
+    if ttype in ttypes.Operator:
+        return value in ["+", "-", "*", "/"]
     return False
 
 
@@ -132,7 +134,8 @@ def insert_implicit_comparisons(tokens, empty_value_map):
         is_compared = (
             next_ttype is ttypes.Comparison or previous_ttype is ttypes.Comparison
         )
-        if token.ttype is ttypes.Name and not is_compared:
+        is_combined = next_ttype is ttypes.Operator or previous_ttype is ttypes.Operator
+        if token.ttype is ttypes.Name and not is_compared and not is_combined:
             column_name = str(token)
             empty_value = empty_value_map[column_name]
             yield sqlparse.sql.Token(ttypes.Punctuation, "(")

tests/test_tpp_backend.py

@@ -830,7 +830,7 @@ def test_patients_satisfying():
         has_asthma=patients.with_these_clinical_events(
             codelist([condition_code], "ctv3")
         ),
-        at_risk=patients.satisfying("(age > 70 AND sex = 'M') OR has_asthma"),
+        at_risk=patients.satisfying("((age + 30) > 100 AND sex = 'M') OR has_asthma"),
     )
     results = study.to_dicts()
     assert [i["at_risk"] for i in results] == ["1", "0", "0", "1"]