Merge pull request #247 from jeromekelleher/num-partitions-named-arg

Change num_partitions to option

Author: jeromekelleher

.github/workflows/ci.yml

@@ -54,14 +54,14 @@ jobs:
           python -m bio2zarr vcf2zarr encode sample.icf sample.vcz -f
       - name: Run distributed explode example
         run: |
-          python -m bio2zarr vcf2zarr dexplode-init tests/data/vcf/sample.vcf.gz sample.icf 3 -f
+          python -m bio2zarr vcf2zarr dexplode-init tests/data/vcf/sample.vcf.gz sample.icf -fn 3
           python -m bio2zarr vcf2zarr dexplode-partition sample.icf 0
           python -m bio2zarr vcf2zarr dexplode-partition sample.icf 1
           python -m bio2zarr vcf2zarr dexplode-partition sample.icf 2
           python -m bio2zarr vcf2zarr dexplode-finalise sample.icf
       - name: Run distributed encode example
         run: |
-          python -m bio2zarr vcf2zarr dencode-init sample.icf sample.vcz 3 -f --variants-chunk-size=3
+          python -m bio2zarr vcf2zarr dencode-init sample.icf sample.vcz -fn 3 --variants-chunk-size=3
           python -m bio2zarr vcf2zarr dencode-partition sample.vcz 0
           python -m bio2zarr vcf2zarr dencode-partition sample.vcz 1
           python -m bio2zarr vcf2zarr dencode-partition sample.vcz 2

CHANGELOG.md

@@ -1,6 +1,9 @@
 # 0.1.0 2024-06-10
 
+- Initial production-ready version.
 - Add -Q/--no-progress flag to CLI
+- Change num-partitions argument in dexplode-init and dencode-init
+  to a named option.
 
 # 0.0.10 2024-05-15
 - Change output format of dexplode-init and dencode-init

bio2zarr/cli.py

@@ -44,7 +44,13 @@ def list_commands(self, ctx):
     "zarr_path", type=click.Path(exists=True, file_okay=False, dir_okay=True)
 )
 
-num_partitions = click.argument("num_partitions", type=click.IntRange(min=1))
+num_partitions = click.option(
+    "-n",
+    "--num-partitions",
+    type=click.IntRange(min=1),
+    default=None,
+    help="Target number of partitions to split into",
+)
 
 partition = click.argument("partition", type=click.IntRange(min=0))
 
@@ -172,6 +178,15 @@ def check_overwrite_dir(path, force):
         shutil.rmtree(tmp_delete_path)
 
 
+def check_partitions(num_partitions):
+    if num_partitions is None:
+        raise click.UsageError(
+            "-n/--num-partitions must currently be specified. Future versions "
+            "will provide reasonable defaults or other means of specifying "
+            "partitions."
+        )
+
+
 def get_compressor(cname):
     if cname is None:
         return None
@@ -249,10 +264,11 @@ def dexplode_init(
 ):
     """
     Initial step for distributed conversion of VCF(s) to intermediate columnar format
-    over the requested number of paritions.
+    over some number of paritions.
     """
     setup_logging(verbose)
     check_overwrite_dir(icf_path, force)
+    check_partitions(num_partitions)
     work_summary = vcf2zarr.explode_init(
         icf_path,
         vcfs,
@@ -388,7 +404,7 @@ def dencode_init(
     """
     Initialise conversion of intermediate format to VCF Zarr. This will
     set up the specified ZARR_PATH to perform this conversion over
-    NUM_PARTITIONS.
+    some number of partitions.
 
     The output of this commmand is the actual number of partitions generated
     (which may be less then the requested number, if there is not sufficient
@@ -400,6 +416,7 @@ def dencode_init(
     """
     setup_logging(verbose)
     check_overwrite_dir(zarr_path, force)
+    check_partitions(num_partitions)
     work_summary = vcf2zarr.encode_init(
         icf_path,
         zarr_path,
@@ -545,21 +562,15 @@ def plink2zarr():
 @version
 @click.argument("vcf_path", type=click.Path(exists=True, dir_okay=False))
 @verbose
-@click.option(
-    "-n",
-    "--num-parts",
-    type=int,
-    default=None,
-    help="Target number of partitions to split the VCF into",
-)
+@num_partitions
 @click.option(
     "-s",
-    "--part-size",
+    "--partition-size",
     type=str,
     default=None,
     help="Target (compressed) size of VCF partitions, e.g. 100KB, 10MiB, 1G.",
 )
-def vcfpartition(vcf_path, verbose, num_parts, part_size):
+def vcfpartition(vcf_path, verbose, num_partitions, partition_size):
     """
     Output bcftools region strings that partition an indexed VCF/BCF file
     into either an approximate number of parts (-n), or parts of approximately
@@ -574,12 +585,14 @@ def vcfpartition(vcf_path, verbose, num_parts, part_size):
     of records that they contain.
     """
     setup_logging(verbose)
-    if num_parts is None and part_size is None:
-        raise click.UsageError("Either --num-parts or --part-size must be specified")
+    if num_partitions is None and partition_size is None:
+        raise click.UsageError(
+            "Either --num-partitions or --partition-size must be specified"
+        )
 
     indexed_vcf = vcf_utils.IndexedVcf(vcf_path)
     regions = indexed_vcf.partition_into_regions(
-        num_parts=num_parts, target_part_size=part_size
+        num_parts=num_partitions, target_part_size=partition_size
     )
     for region in regions:
         click.echo(f"{region}\t{vcf_path}")

docs/vcf2zarr/tutorial.md

@@ -214,7 +214,7 @@ Let's go through the example above using the distributed commands. First, we
 rm -fR sample-dist.icf
 ```
 ```{code-cell}
-vcf2zarr dexplode-init sample.vcf.gz sample-dist.icf -Q 5
+vcf2zarr dexplode-init sample.vcf.gz sample-dist.icf -n 5 -Q
 ```
 
 Here we asked ``dexplode-init`` to set up an ICF store in which the data 

tests/test_cli.py

@@ -116,7 +116,7 @@ def test_vcf_dexplode_init_compressor(self, mocked, tmp_path, compressor):
         runner = ct.CliRunner(mix_stderr=False)
         result = runner.invoke(
             cli.vcf2zarr_main,
-            f"dexplode-init {self.vcf_path} {icf_path} 1 -C {compressor}",
+            f"dexplode-init {self.vcf_path} {icf_path} -n 1 -C {compressor}",
             catch_exceptions=False,
         )
         assert result.exit_code == 0
@@ -294,7 +294,7 @@ def test_vcf_dexplode_init(self, mocked, tmp_path, progress, flag):
         icf_path = tmp_path / "icf"
         result = runner.invoke(
             cli.vcf2zarr_main,
-            f"dexplode-init {self.vcf_path} {icf_path} 5 {flag}",
+            f"dexplode-init {self.vcf_path} {icf_path} -n 5 {flag}",
             catch_exceptions=False,
         )
         assert result.exit_code == 0
@@ -309,7 +309,7 @@ def test_vcf_dexplode_init(self, mocked, tmp_path, progress, flag):
             **args,
         )
 
-    @pytest.mark.parametrize("num_partitions", ["-- -1", "0", "asdf", "1.112"])
+    @pytest.mark.parametrize("num_partitions", ["-1", "0", "asdf", "1.112"])
     @mock.patch("bio2zarr.vcf2zarr.explode_init", return_value=5)
     def test_vcf_dexplode_init_bad_num_partitions(
         self, mocked, tmp_path, num_partitions
@@ -318,11 +318,24 @@ def test_vcf_dexplode_init_bad_num_partitions(
         icf_path = tmp_path / "icf"
         result = runner.invoke(
             cli.vcf2zarr_main,
-            f"dexplode-init {self.vcf_path} {icf_path} {num_partitions}",
+            f"dexplode-init {self.vcf_path} {icf_path} -n {num_partitions}",
             catch_exceptions=False,
         )
         assert result.exit_code == 2
-        assert "Invalid value for 'NUM_PARTITIONS'" in result.stderr
+        assert "Invalid value for '-n'" in result.stderr
+        mocked.assert_not_called()
+
+    @mock.patch("bio2zarr.vcf2zarr.explode_init", return_value=5)
+    def test_vcf_dexplode_init_no_partitions(self, mocked, tmp_path):
+        runner = ct.CliRunner(mix_stderr=False)
+        icf_path = tmp_path / "icf"
+        result = runner.invoke(
+            cli.vcf2zarr_main,
+            f"dexplode-init {self.vcf_path} {icf_path}",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 2
+        assert "-n/--num-partitions must currently be specified" in result.stderr
         mocked.assert_not_called()
 
     @mock.patch("bio2zarr.vcf2zarr.explode_partition")
@@ -457,7 +470,7 @@ def test_dencode_init(self, mocked, tmp_path, progress, flag):
         runner = ct.CliRunner(mix_stderr=False)
         result = runner.invoke(
             cli.vcf2zarr_main,
-            f"dencode-init {icf_path} {zarr_path} 10 {flag}",
+            f"dencode-init {icf_path} {zarr_path} -n 10 {flag}",
             catch_exceptions=False,
         )
         assert result.exit_code == 0
@@ -472,6 +485,21 @@ def test_dencode_init(self, mocked, tmp_path, progress, flag):
             **args,
         )
 
+    @mock.patch("bio2zarr.vcf2zarr.encode_init", return_value=5)
+    def test_vcf_dencode_init_no_partitions(self, mocked, tmp_path):
+        runner = ct.CliRunner(mix_stderr=False)
+        icf_path = tmp_path / "icf"
+        icf_path.mkdir()
+        zarr_path = tmp_path / "zarr"
+        result = runner.invoke(
+            cli.vcf2zarr_main,
+            f"dencode-init {icf_path} {zarr_path}",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 2
+        assert "-n/--num-partitions must currently be specified" in result.stderr
+        mocked.assert_not_called()
+
     @mock.patch("bio2zarr.vcf2zarr.encode_partition")
     def test_vcf_dencode_partition(self, mocked, tmp_path):
         runner = ct.CliRunner(mix_stderr=False)
@@ -601,7 +629,7 @@ def test_dexplode(self, tmp_path, one_based):
         runner = ct.CliRunner(mix_stderr=False)
         result = runner.invoke(
             cli.vcf2zarr_main,
-            f"dexplode-init {self.vcf_path} {icf_path} 5 --json -Q",
+            f"dexplode-init {self.vcf_path} {icf_path} -n 5 --json -Q",
             catch_exceptions=False,
         )
         assert result.exit_code == 0
@@ -678,7 +706,7 @@ def test_dencode(self, tmp_path, one_based):
         assert result.exit_code == 0
         result = runner.invoke(
             cli.vcf2zarr_main,
-            f"dencode-init {icf_path} {zarr_path} 5 --variants-chunk-size=3 --json",
+            f"dencode-init {icf_path} {zarr_path} -n 5 --variants-chunk-size=3 --json",
             catch_exceptions=False,
         )
         assert result.exit_code == 0