Add support for index -n/--nrecords

tests/test_stats.py

@@ -0,0 +1,16 @@
+import pathlib
+from io import StringIO
+
+from vcztools.stats import nrecords
+
+from .utils import vcz_path_cache
+
+
+def test_nrecords():
+    original = pathlib.Path("tests/data/vcf") / "sample.vcf.gz"
+    vcz = vcz_path_cache(original)
+
+    output_str = StringIO()
+    nrecords(vcz, output_str)
+
+    assert output_str.getvalue() == "9\n"

vcztools/cli.py

@@ -2,7 +2,7 @@
 
 import click
 
-from . import regions, vcf_writer
+from . import regions, stats, vcf_writer
 
 
 class NaturalOrderGroup(click.Group):
@@ -16,8 +16,17 @@ def list_commands(self, ctx):
 
 @click.command
 @click.argument("path", type=click.Path())
-def index(path):
-    regions.create_index(path)
+@click.option(
+    "-n",
+    "--nrecords",
+    is_flag=True,
+    help="Print the number of records (variants).",
+)
+def index(path, nrecords):
+    if nrecords:
+        stats.nrecords(path, sys.stdout)
+    else:
+        regions.create_index(path)
 
 
 @click.command

vcztools/stats.py

@@ -0,0 +1,15 @@
+from contextlib import ExitStack
+from pathlib import Path
+
+import zarr
+
+
+def nrecords(vcz, output):
+    root = zarr.open(vcz, mode="r")
+
+    with ExitStack() as stack:
+        if isinstance(output, str) or isinstance(output, Path):
+            output = stack.enter_context(open(output, mode="w"))
+
+        num_variants = root["variant_position"].shape[0]
+        print(num_variants, file=output)