Update query format parser to parse sample loops

sgkit-dev · Oct 2, 2024 · a2ff17c · a2ff17c
1 parent e039823
commit a2ff17c
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 5 deletions.
diff --git a/tests/test_query.py b/tests/test_query.py
@@ -53,6 +53,42 @@ def parser(self):
                 r"Read depth: %INFO/DP\n",
                 ["Read", " ", "depth:", " ", "%INFO/DP", "\n"],
             ),
+            (
+                r"%CHROM\t%POS\t%REF\t%ALT[\t%SAMPLE=%GT]\n",
+                [
+                    "%CHROM",
+                    "\t",
+                    "%POS",
+                    "\t",
+                    "%REF",
+                    "\t",
+                    "%ALT",
+                    ["\t", "%SAMPLE", "=", "%GT"],
+                    "\n",
+                ],
+            ),
+            (
+                r"%CHROM\t%POS\t%REF\t%ALT[\t%SAMPLE=%GT{0}]\n",
+                [
+                    "%CHROM",
+                    "\t",
+                    "%POS",
+                    "\t",
+                    "%REF",
+                    "\t",
+                    "%ALT",
+                    ["\t", "%SAMPLE", "=", ["%GT", 0]],
+                    "\n",
+                ],
+            ),
+            (
+                r"GQ:[ %GQ] \t GT:[ %GT]\n",
+                ["GQ:", [" ", "%GQ"], " ", "\t", " ", "GT:", [" ", "%GT"], "\n"],
+            ),
+            (
+                r"[%SAMPLE %GT %DP\n]",
+                [["%SAMPLE", " ", "%GT", " ", "%DP", "\n"]],
+            ),
         ],
     )
     def test_valid_expressions(self, parser, expression, expected_result):

diff --git a/vcztools/query.py b/vcztools/query.py
@@ -35,17 +35,21 @@ def __init__(self):
         )
         newline_pattern = pp.Literal("\\n").set_parse_action(pp.replace_with("\n"))
         tab_pattern = pp.Literal("\\t").set_parse_action(pp.replace_with("\t"))
-        pattern = pp.ZeroOrMore(
-            subscript_pattern
+        format_pattern = pp.Forward()
+        sample_loop_pattern = pp.Group(
+            pp.Literal("[").suppress() + format_pattern + pp.Literal("]").suppress()
+        )
+        format_pattern <<= pp.ZeroOrMore(
+            sample_loop_pattern
+            | subscript_pattern
             | tag_pattern
             | newline_pattern
             | tab_pattern
             | pp.White()
             | pp.Word(pp.printables, exclude_chars=r"\{}[]%")
-        )
-        pattern = pattern.leave_whitespace()
+        ).leave_whitespace()
 
-        self._parser = functools.partial(pattern.parse_string, parse_all=True)
+        self._parser = functools.partial(format_pattern.parse_string, parse_all=True)
 
     def __call__(self, *args, **kwargs):
         assert len(args) == 1