Add query list samples command

tests/test_bcftools_validation.py

@@ -52,7 +52,7 @@ def run_vcztools(args: str) -> str:
     ]
 )
 # fmt: on
-def test(tmp_path, args, vcf_file):
+def test_vcf_output(tmp_path, args, vcf_file):
     original = pathlib.Path("tests/data/vcf") / vcf_file
     vcz = vcz_path_cache(original)
 
@@ -67,3 +67,20 @@ def test(tmp_path, args, vcf_file):
         f.write(vcztools_out)
 
     assert_vcfs_close(bcftools_out_file, vcztools_out_file)
+
+
+@pytest.mark.parametrize(
+    ("args", "vcf_name"),
+    [
+        ("query -l", "sample.vcf.gz"),
+        ("query --list-samples", "1kg_2020_chrM.vcf.gz"),
+    ],
+)
+def test_output(tmp_path, args, vcf_name):
+    vcf_path = pathlib.Path("tests/data/vcf") / vcf_name
+    vcz_path = vcz_path_cache(vcf_path)
+
+    bcftools_output = run_bcftools(f"{args} {vcf_path}")
+    vcztools_output = run_vcztools(f"{args} {vcz_path}")
+
+    assert vcztools_output == bcftools_output

tests/test_query.py

@@ -0,0 +1,20 @@
+import pathlib
+from io import StringIO
+
+from tests.utils import vcz_path_cache
+from vcztools.query import list_samples
+
+
+def test_list_samples(tmp_path):
+    vcf_path = pathlib.Path("tests/data/vcf") / "sample.vcf.gz"
+    vcz_path = vcz_path_cache(vcf_path)
+    expected_output = "NA00001\nNA00002\nNA00003\n"
+
+    with StringIO() as output:
+        list_samples(vcz_path, output)
+        assert output.getvalue() == expected_output
+
+    list_samples(vcz_path, output=tmp_path / "sample_ids.txt")
+
+    with open(tmp_path / "sample_ids.txt") as file:
+        assert file.read() == expected_output

vcztools/cli.py

@@ -2,6 +2,7 @@
 
 import click
 
+from . import query as query_module
 from . import regions, vcf_writer
 
 
@@ -20,6 +21,20 @@ def index(path):
     regions.create_index(path)
 
 
+@click.command
+@click.argument("path", type=click.Path())
+@click.option(
+    "-l",
+    "--list-samples",
+    is_flag=True,
+    help="List the sample IDs and exit.",
+)
+def query(path, list_samples):
+    if list_samples:
+        query_module.list_samples(path)
+        return
+
+
 @click.command
 @click.argument("path", type=click.Path())
 @click.option(
@@ -96,4 +111,5 @@ def vcztools_main():
 
 
 vcztools_main.add_command(index)
+vcztools_main.add_command(query)
 vcztools_main.add_command(view)

vcztools/query.py

@@ -0,0 +1,15 @@
+import pathlib
+from contextlib import ExitStack
+
+import zarr
+
+
+def list_samples(vcz_path, output=None):
+    root = zarr.open(vcz_path, mode="r")
+
+    with ExitStack() as stack:
+        if isinstance(output, str) or isinstance(output, pathlib.Path):
+            output = stack.enter_context(open(output, mode="w"))
+
+        sample_ids = root["sample_id"][:]
+        print("\n".join(sample_ids), file=output)