Add unary minus

More tests for arithmetic

Author: jeromekelleher

tests/test_bcftools_validation.py

@@ -162,15 +162,6 @@ def test_vcf_output_with_output_option(tmp_path, args, vcf_file):
             r"query -f '[%POS %QUAL\n]' -i'(QUAL > 10 && POS > 100000)'",
             "sample.vcf.gz",
         ),
-        # Check arithmetic evaluation in filter queries. All these should
-        # result to POS=112, which exists.
-        (r"query -f '%POS\n' -i 'POS=(112 + 1)'", "sample.vcf.gz"),
-        (r"query -f '%POS\n' -i 'POS =(224 / 2)'", "sample.vcf.gz"),
-        (r"query -f '%POS\n' -i 'POS= (112 * 3) / 3'", "sample.vcf.gz"),
-        (r"query -f '%POS\n' -i 'POS=(112 * 3 / 3   )'", "sample.vcf.gz"),
-        (r"query -f '%POS\n' -i 'POS=25 * 4 + 24 / 2'", "sample.vcf.gz"),
-        (r"query -f '%POS\n' -i 'POS=112 * -1 * -1'", "sample.vcf.gz"),
-        (r"query -f '%POS\n' -i '-POS=-112'", "sample.vcf.gz"),
     ],
 )
 def test_output(tmp_path, args, vcf_name):
@@ -194,7 +185,9 @@ def test_output(tmp_path, args, vcf_name):
         "POS=(112 * 3 / 3   )",
         "POS=25 * 4 + 24 / 2",
         "POS=112 * -1 * -1",
-        # "-POS=-112",
+        "-POS=-112",
+        "POS=112.25 - 1 / 4",
+        "POS=112.25e3 * 1e-3 - 0.25",
     ],
 )
 def test_query_arithmethic(tmp_path, expr):
@@ -211,7 +204,6 @@ def test_query_arithmethic(tmp_path, expr):
     assert vcztools_output == "112\n"
 
 
-
 @pytest.mark.parametrize(
     ("args", "vcf_name"),
     [

tests/test_filter.py

@@ -92,6 +92,7 @@ def test_evaluate(self, expression, data, expected):
             ("a == b", {"variant_a", "variant_b"}),
             ("a == b + c", {"variant_a", "variant_b", "variant_c"}),
             ("(a + 1) < (b + c) - d / a", {f"variant_{x}" for x in "abcd"}),
+            ("-(a + b)", {f"variant_{x}" for x in "ab"}),
         ],
     )
     def test_referenced_fields(self, expr, expected):
@@ -103,6 +104,7 @@ def test_referenced_fields(self, expr, expected):
         [
             ("a == b", "(variant_a)==(variant_b)"),
             ("a + 1", "(variant_a)+(1)"),
+            ("-a + 1", "(-(variant_a))+(1)"),
             ("a + 1 + 2", "(variant_a)+(1)+(2)"),
             ("a + (1 + 2)", "(variant_a)+((1)+(2))"),
             ("POS<10", "(variant_position)<(10)"),
@@ -139,6 +141,7 @@ class TestBcftoolsParser:
             "(1 + 2) == (1 + 2 + 3)",
             "(1 == 1) != (2 == 2)",
             '("x" == "x")',
+            "-1 == 1 + 2 - 4",
         ],
     )
     def test_python_arithmetic_expressions(self, expr):
@@ -157,6 +160,8 @@ def test_python_arithmetic_expressions(self, expr):
             ("(a - b) < (a + b)", {"a": 7, "b": 6}),
             ("(a - b) < (a + b)", {"a": 7.0, "b": 6.666}),
             ("a == a", {"a": 1}),
+            ("-a == -a", {"a": 1}),
+            ("-a == b", {"a": 1, "b": -1}),
             ('a == "string"', {"a": "string"}),
         ],
     )

vcztools/filter.py

@@ -9,6 +9,9 @@
 
 logger = logging.getLogger(__name__)
 
+# Parsing is WAY slower without this!
+pp.ParserElement.enablePackrat()
+
 
 # The parser and evaluation model here are based on the eval_arith example
 # in the pyparsing docs:
@@ -52,6 +55,27 @@ def referenced_fields(self):
         return frozenset([self.field_name])
 
 
+# NOTE we should perhaps add a Operator superclass of UnaryMinus,
+# BinaryOperator and ComparisonOperator to reduce duplication
+# when doing things like referenced_fields. We should probably
+# be extracting the operators and operands once.
+
+
+class UnaryMinus(EvaluationNode):
+    def eval(self, data):
+        op, operand = self.tokens
+        assert op == "-"
+        return -1 * operand.eval(data)
+
+    def __repr__(self):
+        _, operand = self.tokens
+        return f"-({repr(operand)})"
+
+    def referenced_fields(self):
+        _, operand = self.tokens
+        return operand.referenced_fields()
+
+
 class BinaryOperator(EvaluationNode):
     op_map = {
         "*": operator.mul,
@@ -143,8 +167,7 @@ def make_bcftools_filter_parser(all_fields=None, map_vcf_identifiers=True):
     filter_expression = pp.infix_notation(
         constant | identifier,
         [
-            # FIXME Does bcftools support unary minus?
-            # ("-", 1, pp.OpAssoc.RIGHT, ),
+            ("-", 1, pp.OpAssoc.RIGHT, UnaryMinus),
             (pp.one_of("* /"), 2, pp.OpAssoc.LEFT, BinaryOperator),
             (pp.one_of("+ -"), 2, pp.OpAssoc.LEFT, BinaryOperator),
             (comp_op, 2, pp.OpAssoc.LEFT, ComparisonOperator),