diff --git a/tests/test_query.py b/tests/test_query.py index 97f5bce..91b62a8 100644 --- 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 index a1e40eb..6ad0a09 100644 --- 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