Fix KeyError for contigs missing from assembly in write_viral_gff

Contigs present in annotation TSVs but absent from the assembly FASTA
caused a KeyError in three places in write_gff()

This happens when the user provide proteins, as the proteins could be
for contigs that were filtered out by the length filter VIRIfy use.
I've added a bit of code to make it fail if all the contigs are
discarded as this would a problem.

Adds a regression test with some file fixtures.

Author: mberacochea

bin/write_viral_gff.py

@@ -148,9 +148,7 @@ def aggregate_annotations(
                     # and prophage predictions can extend beyond the original contig boundaries
                     clean_contig_name = Record.remove_prophage_from_contig(contig)
                     contig_len = contigs_len_dict[clean_contig_name]
-                    does_the_prophage_overrun = (
-                        prophage_end > contig_len
-                    )
+                    does_the_prophage_overrun = prophage_end > contig_len
 
                     if does_the_prophage_overrun:
                         # We truncate as the prophage_end could overrun
@@ -177,7 +175,7 @@ def aggregate_annotations(
                 if best_hit != "No hit":
                     best_hit = best_hit.replace(".faa", "")
                     viphog_annotation = ";".join(
-                        [f"viphog={best_hit}", f'viphog_taxonomy={row["Label"]}']
+                        [f"viphog={best_hit}", f"viphog_taxonomy={row['Label']}"]
                     )
                     # We need to remove all the virify prophage annotations, if any
                     contig_name_clean = Record.remove_prophage_from_contig(contig)
@@ -212,6 +210,7 @@ def write_gff(
     virify_quality,
     contigs_len_dict,
     ena_mapping=None,
+    user_proteins=False,
 ):
     """Generate a GFF3 file from VIRify output files with comprehensive viral sequence annotations.
 
@@ -229,7 +228,8 @@ def write_gff(
     :param ena_mapping: Optional ENA contig mapping for renaming (ERZ accession will be used if provided)
     :param contigs_len_dict: Optional pre-loaded dictionary mapping contig names to lengths.
                             If not provided, will be loaded from assembly_file.
-
+    :param ena_mapping: ENA mapping dict
+    :param user_proteins: Flag used when users provide their "proteins"
     :return: None (writes GFF file to disk)
     """
     if ena_mapping:
@@ -307,13 +307,33 @@ def empty_if_number(string):
     # Collect all sequence-region headers
     sequence_regions = []
     used_contigs = set()
+
+    missed_contigs = 0
+
     for contig_name in viral_sequences.keys():
         clean_contig_name = Record.remove_prophage_from_contig(contig_name)
         if clean_contig_name not in used_contigs:
             used_contigs.add(clean_contig_name)
-            contig_length = contigs_len_dict[clean_contig_name]
+            # Users may provide proteins for all the contigs, but VIRify only considers contigs
+            # that are longer than 150K, so when users provide a proteins file (--user_proteins)
+            # we allow mismatches here. Guard applies regardless of user_proteins to avoid
+            # writing None as contig length in the GFF3 sequence-region directive.
+            contig_length = contigs_len_dict.get(clean_contig_name)
+            if contig_length is None:
+                missed_contigs += 1
+                continue
             sequence_regions.append((clean_contig_name, contig_length))
 
+    if missed_contigs > 0:
+        logging.warning(
+            f"{missed_contigs} contigs were not found in the assembly and were skipped"
+        )
+
+    if not sequence_regions:
+        raise ValueError(
+            "All the contigs that came from the annotated viral sequences were discarded."
+        )
+
     # Sort sequence-region headers by contig name
     sequence_regions.sort(key=lambda x: x[0])
 
@@ -331,7 +351,9 @@ def empty_if_number(string):
             element_category = "viral_sequence"
             id_ = f"ID={clean_contig_name}|viral_sequence"
             start = 1
-            end = contigs_len_dict[clean_contig_name]
+            end = contigs_len_dict.get(clean_contig_name)
+            if end is None:
+                continue
             mobile_element_type = viral_seq_type
 
             if "prophage" in viral_seq_type:
@@ -385,9 +407,12 @@ def empty_if_number(string):
             region_name = "_".join(cds_id.split("_")[:-1])
             cds_id = cds_id.replace("prophage-0:", "prophage-1:")
 
-            # TODO: review this rule.
-            if end > contigs_len_dict[contig_name]:
-                end = contigs_len_dict[contig_name]
+            contig_len = contigs_len_dict.get(contig_name)
+            if contig_len is None:
+                continue
+
+            if end > contig_len:
+                end = contig_len
 
             quality = (
                 virify_quality[region_name]
@@ -570,6 +595,7 @@ def empty_if_number(string):
     )
 
     logging.info("Generating the gff output")
+
     write_gff(
         checkv_files,
         taxonomy_files,
@@ -580,4 +606,5 @@ def empty_if_number(string):
         virify_quality,
         contigs_len_dict,
         ena_mapping=ena_mapping,
+        user_proteins=args.use_proteins,
     )

modules/local/write_gff/main.nf

@@ -13,13 +13,13 @@ process WRITE_GFF {
     script:
     def use_proteins_flag = params.use_proteins ? "--use-proteins": "" ;
     """
-    write_viral_gff.py \
-    $use_proteins_flag \
-    -v ${viphos_annotations.join(' ')} \
-    -c ${quality_summaries.join(' ')} \
-    -t ${taxonomies.join(' ')} \
-    -s ${meta.id} \
-    -a ${fasta}
+    write_viral_gff.py \\
+      $use_proteins_flag \\
+      -v ${viphos_annotations.join(' ')} \\
+      -c ${quality_summaries.join(' ')} \\
+      -t ${taxonomies.join(' ')} \\
+      -s ${meta.id} \\
+      -a ${fasta}
 
     gt gff3validator ${meta.id}_virify.gff
     """

subworkflows/local/annotate.nf

@@ -14,10 +14,6 @@ include { TABIX_BGZIP                 } from '../../modules/nf-core/tabix/bgzip/
 include { TABIX_BGZIPTABIX            } from '../../modules/nf-core/tabix/bgziptabix/main'
 
 /* Local modules */
-include { VIRSORTER                   } from '../../modules/local/virsorter' 
-include { VIRFINDER                   } from '../../modules/local/virfinder' 
-include { PPRMETA                     } from '../../modules/local/pprmeta'
-include { LENGTH_FILTERING            } from '../../modules/local/length_filtering'  
 include { RATIO_EVALUE                } from '../../modules/local/ratio_evalue' 
 include { ANNOTATION                  } from '../../modules/local/annotation' 
 include { ASSIGN                      } from '../../modules/local/assign' 
@@ -54,6 +50,7 @@ workflow ANNOTATE {
 
     // prodigal
     PREDICT_PROTEINS( input_fastas )
+
     contigs = PREDICT_PROTEINS.out.contigs
     proteins = PREDICT_PROTEINS.out.proteins
     predicted_contigs = PREDICT_PROTEINS.out.predicted_contigs
@@ -88,12 +85,12 @@ workflow ANNOTATE {
       checkv_db.first()
     )
 
-    viphos_annotations = ANNOTATION.out.annotations.map{meta, type, annotation -> [meta, annotation]}.groupTuple()
-    taxonomy_annotations = ASSIGN.out.map{meta, type, annotation -> [meta, annotation]}.groupTuple()
-    checkv_results = CHECKV.out.map{meta, type, quality -> [meta, quality]}.groupTuple()
+    viphos_annotations = ANNOTATION.out.annotations.map { meta, _type, annotation -> [meta, annotation] }.groupTuple()
+    taxonomy_annotations = ASSIGN.out.map { meta, _type, annotation -> [meta, annotation] }.groupTuple()
+    checkv_results = CHECKV.out.map { meta, _type, quality -> [meta, quality] }.groupTuple()
 
     WRITE_GFF(
-      contigs.join(viphos_annotations).join(taxonomy_annotations).join(checkv_results)
+      contigs.join(viphos_annotations).join(taxonomy_annotations).join(checkv_results),
     )
 
     /**********************************************/

tests/test_write_gff.py

@@ -1,7 +1,9 @@
 #!/bin/env python3
 
 import os
+from pathlib import Path
 import unittest
+from unittest.mock import patch
 import glob
 import hashlib
 
@@ -167,3 +169,56 @@ def test_prophage_coordinate_truncation_in_gff(self):
         # Clean up
         if os.path.exists("test_sample_virify.gff"):
             os.unlink("test_sample_virify.gff")
+
+
+def test_contig_missing_from_assembly_is_skipped_with_warning(tmp_path):
+    """Contig present in annotation TSV but absent from the assembly FASTA is skipped.
+
+    Regression test for KeyError: '<contig_name>' when a contig in the
+    annotation file has no matching entry in contigs_len_dict.
+    Both the sequence-region header, the mobile-element record, and any
+    associated CDS records for the missing contig must be omitted from the
+    output GFF, and a WARNING must be emitted reporting the count.
+    """
+
+    fixtures = Path(__file__).parent / "write_gff_missing_contig_fixtures"
+    assembly_fasta = fixtures / "assembly.fasta"
+    annotation_tsv = fixtures / "high_confidence_viral_contigs_annotation.tsv"
+    checkv_tsv = fixtures / "high_confidence_viral_contigs_quality_summary.tsv"
+    taxonomy_tsv = fixtures / "high_confidence_viral_contigs_annotation_taxonomy.tsv"
+
+    contigs_len_dict = get_contig_lengths_per_contig(str(assembly_fasta))
+
+    viral_sequences, cds_annotations, virify_quality = aggregate_annotations(
+        [str(annotation_tsv)], contigs_len_dict
+    )
+
+    # Both contigs appear in annotation output before write_gff filtering
+    assert "valid_contig" in viral_sequences
+    assert "missing_contig" in viral_sequences
+
+    # Passing a full path as sample_prefix directs the output GFF into tmp_path
+    sample_prefix = str(tmp_path / "test_missing_contig")
+
+    with patch("logging.warning") as mock_warn:
+        write_gff(
+            [str(checkv_tsv)],
+            [str(taxonomy_tsv)],
+            sample_prefix,
+            str(assembly_fasta),
+            viral_sequences,
+            cds_annotations,
+            virify_quality,
+            contigs_len_dict,
+        )
+
+    warned_messages = [str(call) for call in mock_warn.call_args_list]
+    assert any(
+        "1 contigs were not found in the assembly and were skipped" in msg
+        for msg in warned_messages
+    )
+
+    output_gff = tmp_path / "test_missing_contig_virify.gff"
+    gff_content = output_gff.read_text()
+    assert "valid_contig" in gff_content
+    assert "missing_contig" not in gff_content

tests/write_gff_missing_contig_fixtures/assembly.fasta

@@ -0,0 +1,16 @@
+>valid_contig
+GCTACCTGGTGCACCGTGCCGCCTTGCCAGCCCAATGAGTGCATGAGATGTAAAGGAGTATTTGTCATGG
+TATCGCCCTTTCTTTGATGTGCGCCCTGGTTGGCGCGGTAGAAAGATAATAGCCGTTTCCTTTATAGATA
+GCAAGCAAAAAAGAAATGAAAAAAAGATTAAAAAAAATGGTTGACGTTCCTCACCTGGTAGCGTAGTATT
+CATCTTGTCGAGGGGCGGCAATAGTGCCGCTCATAACAGAAGAGAGAAATAATCATGGCATCATATAGCG
+CATCAATGATCGTAAAGGGTATCACTGGTAAACCTATGCGCGGTGGCGGATCTTTCGCTAACATCAATAT
+TGATGGACGTTTATCCATGGACGCAGCGATCACAGTGGCACGGGAAACATTTAAAAAAGAATGTAATTTC
+AATAAACAAGATTATCTCGGTTTCGCTATTGAAAAAACCGCGCGCTTTGTGGATTATAAAAGCCCGAAAA
+TGATCGACACTACACTAAAAGCGAAAGATGTAGCGTTTTTGCTTTAAGGCTTGACACTGAAAAATATTTT
+GGTAATATATACCACATAAGGCGGGGAAATGGTTTCCCGCCAAATCAAAGAAAGGGTAATGCAATGAGCA
+AGCAAATCAATACTATCATGTTTAATCTGGTATCTTTCCGCGATGACGTGAAAAACCTTCCGCGTGATGC
+AGTGGATGCACGGATCGAAAGCTATCGCACACAGATCGCAACGTTACCGCTGAAACGTGATCAACATGCC
+GCTAATATGATGCTTGATGCTATGATTAAAAAGATGATTCAATGTGATCTTACCCTGGCTTATGAATATG
+CTGGCGATCTCTTTGTAACATATAGCAAGCCCGTTCCAGGTATGAGCACCACGGAAGAATTGCATGCCTT
+AAATCTGCGACAAGAAAATCTGCGACAAGAAAATCTGCGACAAGAAAATCTGCGACAAGAAAATCTGCGA
+CAAGAAAATCTGCGACAAGA

tests/write_gff_missing_contig_fixtures/high_confidence_viral_contigs_annotation.tsv

@@ -0,0 +1,3 @@
+Contig	CDS_ID	Start	End	Direction	Best_hit	Abs_Evalue_exp	Label
+valid_contig	valid_contig_1	1	100	1	No hit	NA
+missing_contig	missing_contig_1	1	100	1	No hit	NA

tests/write_gff_missing_contig_fixtures/high_confidence_viral_contigs_annotation_taxonomy.tsv

@@ -0,0 +1,3 @@
+contig_ID	superkingdom	kingdom	phylum	subphylum	class	order	suborder	family	subfamily	genus
+valid_contig										
+missing_contig										

tests/write_gff_missing_contig_fixtures/high_confidence_viral_contigs_quality_summary.tsv

@@ -0,0 +1,3 @@
+contig_id	contig_length	provirus	proviral_length	gene_count	viral_genes	host_genes	checkv_quality	miuvig_quality	completeness	completeness_method	contamination	kmer_freq	warnings
+valid_contig	1000	No	NA	2	1	0	Medium-quality	Genome-fragment	80.0	HMM-based	0.0	1.0
+missing_contig	1000	No	NA	2	1	0	Medium-quality	Genome-fragment	80.0	HMM-based	0.0	1.0