Repository: griffithlab/VAtools

Issue:
Add feature to vcf-genotype-annotator to allow for adding GT field to existing samples
Right now the `vcf-genotype-annotator` will throw a fatal error if the give sample already exists in the VCF. If the GT field does not yet exist in the VCF (no header and not a field in `FORMAT`) this new option would allow for the tool to add the GT field.

Comments:

Patch:
diff --git a/docs/index.rst b/docs/index.rst
index 7f26fc7..3c991e0 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -24,7 +24,8 @@ annotate VCF files with data from other tools.
     field in the VCF INFO column.
 
 **vcf-genotype-annotator**
-    A tool to add a new sample to an existing VCF file.
+    A tool to add a new sample to an existing VCF file or fill in the GT field
+    for an existing sample in a VCF.
 
 **vep-annotation-reporter**
     A tool to create a tab-delimited (TSV) file of variants in a VCF and their
diff --git a/docs/vcf_genotype_annotator.rst b/docs/vcf_genotype_annotator.rst
index 4781da6..5948461 100644
--- a/docs/vcf_genotype_annotator.rst
+++ b/docs/vcf_genotype_annotator.rst
@@ -7,6 +7,9 @@ The sample's GT field is pre-populated with a default value given by
 the third positional argument. Options are ``0/1``, ``1/1``, ``0/0``, or
 ``.``.
 
+It can also be used to add a GT field to an existing sample, e.g. for VCFs
+created by Strelka which does not output a GT field for its calls.
+
 By default the output VCF will be written to a ``.genotype.vcf`` file next to
 your input VCF file. You can set a different output file using the
 ``--output-vcf`` parameter.
diff --git a/vatools/vcf_genotype_annotator.py b/vatools/vcf_genotype_annotator.py
index 7df8942..65f7969 100644
--- a/vatools/vcf_genotype_annotator.py
+++ b/vatools/vcf_genotype_annotator.py
@@ -7,8 +7,9 @@ from collections import OrderedDict
 def create_vcf_reader(args):
     vcf_reader = vcfpy.Reader.from_path(args.input_vcf)
     if args.sample_name in vcf_reader.header.samples.names:
-        vcf_reader.close()
-        raise Exception("VCF already contains a sample column for sample {}.".format(args.sample_name))
+        if 'GT' in vcf_reader.header.format_ids():
+            vcf_reader.close()
+            raise Exception("VCF already contains a sample column for sample {} with a GT field.".format(args.sample_name))
     return vcf_reader
 
 def create_vcf_writer(args, vcf_reader):
@@ -18,14 +19,18 @@ def create_vcf_writer(args, vcf_reader):
         (head, sep, tail) = args.input_vcf.rpartition('.vcf')
         output_file = ('').join([head, '.genotype.vcf', tail])
     sample_info = vcf_reader.header.samples
-    sample_info.names.append(args.sample_name)
-    sample_info.name_to_idx[args.sample_name] = len(sample_info.names)-1
+    if args.sample_name in sample_info.names:
+        append_to_existing_sample = True
+    else:
+        append_to_existing_sample = False
+        sample_info.names.append(args.sample_name)
+        sample_info.name_to_idx[args.sample_name] = len(sample_info.names)-1
     new_header = vcfpy.Header(samples = sample_info)
     for line in vcf_reader.header.lines:
         if not (line.key == 'FORMAT' and line.id == 'GT'):
             new_header.add_line(line)
     new_header.add_format_line(OrderedDict([('ID', 'GT'), ('Number', '1'), ('Type', 'String'), ('Description', 'Genotype')]))
-    return vcfpy.Writer.from_path(output_file, new_header)
+    return ( vcfpy.Writer.from_path(output_file, new_header), append_to_existing_sample )
 
 def define_parser():
     parser = argparse.ArgumentParser("vcf-genotype-annotator")
@@ -56,20 +61,23 @@ def main(args_input = sys.argv[1:]):
     args = parser.parse_args(args_input)
 
     vcf_reader = create_vcf_reader(args)
-    vcf_writer = create_vcf_writer(args, vcf_reader)
+    (vcf_writer, append_to_existing_sample) = create_vcf_writer(args, vcf_reader)
 
     for entry in vcf_reader:
-        new_sample_call = vcfpy.Call(args.sample_name, data={'GT': args.genotype_value})
         if "GT" not in entry.FORMAT:
             if isinstance(entry.FORMAT, tuple):
                 entry.FORMAT = ["GT"]
             else:
                 entry.FORMAT.insert(0, 'GT')
-        if entry.calls:
-            entry.calls.append(new_sample_call)
+        if append_to_existing_sample:
+            entry.call_for_sample[args.sample_name].data['GT'] = args.genotype_value
         else:
-            entry.calls = [new_sample_call]
-        entry.call_for_sample = {call.sample: call for call in entry.calls}
+            new_sample_call = vcfpy.Call(args.sample_name, data={'GT': args.genotype_value})
+            if entry.calls:
+                entry.calls.append(new_sample_call)
+            else:
+                entry.calls = [new_sample_call]
+            entry.call_for_sample = {call.sample: call for call in entry.calls}
         vcf_writer.write_record(entry)
 
     vcf_reader.close()


Test patch:
diff --git a/tests/test_vcf_genotype_annotator.py b/tests/test_vcf_genotype_annotator.py
index 8343771..28a1322 100644
--- a/tests/test_vcf_genotype_annotator.py
+++ b/tests/test_vcf_genotype_annotator.py
@@ -16,7 +16,7 @@ class VcfExpressionEncoderTests(unittest.TestCase):
     def test_source_compiles(self):
         self.assertTrue(py_compile.compile(self.executable))
 
-    def test_error_sample_name_already_exists(self):
+    def test_error_sample_name_already_exists_with_GT_field(self):
         with self.assertRaises(Exception) as context:
             command = [
                 os.path.join(self.test_data_dir, 'input.vcf'),
@@ -24,7 +24,7 @@ class VcfExpressionEncoderTests(unittest.TestCase):
                 '0/1',
             ]
             vcf_genotype_annotator.main(command)
-        self.assertTrue('VCF already contains a sample column for sample H_NJ-HCC1395-HCC1395.' in str(context.exception))
+        self.assertTrue('VCF already contains a sample column for sample H_NJ-HCC1395-HCC1395 with a GT field.' in str(context.exception))
 
     def test_no_sample_vcf(self):
         temp_path = tempfile.TemporaryDirectory()
@@ -61,3 +61,15 @@ class VcfExpressionEncoderTests(unittest.TestCase):
         vcf_genotype_annotator.main(command)
         self.assertTrue(cmp(os.path.join(self.test_data_dir, 'no_gt_in_format.genotype.vcf'), os.path.join(temp_path.name, 'input.genotype.vcf')))
         temp_path.cleanup()
+
+    def test_adding_gt_in_existing_sample(self):
+        temp_path = tempfile.TemporaryDirectory()
+        os.symlink(os.path.join(self.test_data_dir, 'input.no_gt_in_format.vcf'), os.path.join(temp_path.name, 'input.vcf'))
+        command = [
+            os.path.join(temp_path.name, 'input.vcf'),
+            'H_NJ-HCC1395-HCC1395',
+            '0/1',
+        ]
+        vcf_genotype_annotator.main(command)
+        self.assertTrue(cmp(os.path.join(self.test_data_dir, 'existing_sample.genotype.vcf'), os.path.join(temp_path.name, 'input.genotype.vcf')))
+        temp_path.cleanup()