Add contig-id argument

Author: szhan

python/tests/test_cli.py

@@ -263,29 +263,37 @@ def test_fasta_long_args(self):
         assert args.tree_sequence == tree_sequence
         assert args.wrap == 50
 
-    def test_vcf_default_values(self):
+    @pytest.mark.parametrize(
+        "flags,expected",
+        (
+            [[], None],
+            [["-P", "2"], 2],
+            [["--ploidy", "5"], 5],
+        ),
+    )
+    def test_vcf_ploidy(self, flags, expected):
         parser = cli.get_tskit_parser()
         cmd = "vcf"
         tree_sequence = "test.trees"
-        args = parser.parse_args([cmd, tree_sequence])
+        args = parser.parse_args([cmd, tree_sequence, *flags])
         assert args.tree_sequence == tree_sequence
-        assert args.ploidy is None
-
-    def test_vcf_short_args(self):
-        parser = cli.get_tskit_parser()
-        cmd = "vcf"
-        tree_sequence = "test.trees"
-        args = parser.parse_args([cmd, tree_sequence, "-P", "2"])
-        assert args.tree_sequence == tree_sequence
-        assert args.ploidy == 2
-
-    def test_vcf_long_args(self):
+        assert args.ploidy == expected
+
+    @pytest.mark.parametrize(
+        "flags,expected",
+        (
+            [[], "1"],
+            [["-c", "chrX"], "chrX"],
+            [["--contig-id", "chr20"], "chr20"],
+        ),
+    )
+    def test_vcf_contig_id(self, flags, expected):
         parser = cli.get_tskit_parser()
         cmd = "vcf"
         tree_sequence = "test.trees"
-        args = parser.parse_args([cmd, tree_sequence, "--ploidy", "5"])
+        args = parser.parse_args([cmd, tree_sequence, *flags])
         assert args.tree_sequence == tree_sequence
-        assert args.ploidy == 5
+        assert args.contig_id == expected
 
     def test_upgrade_default_values(self):
         parser = cli.get_tskit_parser()

python/tskit/cli.py

@@ -1,7 +1,7 @@
 #
 # MIT License
 #
-# Copyright (c) 2018-2021 Tskit Developers
+# Copyright (c) 2018-2022 Tskit Developers
 # Copyright (c) 2015-2018 University of Oxford
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -132,7 +132,7 @@ def run_fasta(args):
 
 def run_vcf(args):
     tree_sequence = load_tree_sequence(args.tree_sequence)
-    tree_sequence.write_vcf(sys.stdout, ploidy=args.ploidy)
+    tree_sequence.write_vcf(sys.stdout, ploidy=args.ploidy, contig_id=args.contig_id)
 
 
 def add_tree_sequence_argument(parser):
@@ -215,6 +215,9 @@ def get_tskit_parser():
             "individuals"
         ),
     )
+    parser.add_argument(
+        "--contig-id", "-c", type=str, default="1", help="Specify the contig id"
+    )
     parser.set_defaults(runner=run_vcf)
 
     parser = subparsers.add_parser(