ENH: add action to collate genomes to GenomeData[DNASequences] (#99)

DorielaGrabocka · ebolyen · misialq · web-flow · commit 2dbb69f9e919 · 2024-09-20T16:36:00.000+02:00
* ENH: action to convert FeatureData to Genome data added * ENH: action to convert FeatureData to Genome data refined * ENH: action to convert FeatureData to Genome data tests added * ENH: set up test paths added * ENH: imports modified * REF: fix private imports for q2-types refactor (#95) * ENH: changes requested done * ENH: imports modified * ENH: description modified * Update q2_assembly/plugin_setup.py * ENH: tests modified * ENH: tests modified * ENH: checking for duplicate IDs code added * ENH: accounting for duplicates case * ENH: accounting for duplicates case, modifying tests to check file content * ENH: code refactoring * ENH: tests corrected * Update q2_assembly/helpers/helpers.py --------- Co-authored-by: Evan Bolyen <ebolyen@gmail.com> Co-authored-by: Michal Ziemski <mziemski@ethz.ch>
diff --git a/q2_assembly/helpers/helpers.py b/q2_assembly/helpers/helpers.py
@@ -9,9 +9,14 @@
 import os
 import shutil
 import warnings
+from typing import Union
+from warnings import warn
 
 import numpy as np
+import skbio.io
 from q2_types.bowtie2 import Bowtie2IndexDirFmt
+from q2_types.feature_data import DNAFASTAFormat, DNAIterator
+from q2_types.genome_data import GenomeSequencesDirectoryFormat
 from q2_types.per_sample_sequences import BAMDirFmt, ContigSequencesDirFmt
 from qiime2.util import duplicate
 
@@ -109,3 +114,50 @@ def collate_alignments(alignments: BAMDirFmt) -> BAMDirFmt:
             duplicate(_alignment, os.path.join(collated_alignments.path, filename))
 
     return collated_alignments
+
+
+def collate_genomes(
+    genomes: Union[DNAFASTAFormat, GenomeSequencesDirectoryFormat],
+    on_duplicates: str = "warn",
+) -> GenomeSequencesDirectoryFormat:
+    genomes_dir = GenomeSequencesDirectoryFormat()
+    error_on_duplicates = True if on_duplicates == "error" else False
+    ids = set()
+    duplicate_ids = set()
+    msg = "Duplicate sequence files were found for the following IDs: {}."
+    if isinstance(genomes[0], DNAFASTAFormat):
+        for genome_file in genomes:
+            for genome in genome_file.view(DNAIterator):
+                fn = genome.metadata["id"]
+                if fn not in ids:
+                    with open(os.path.join(genomes_dir.path, fn + ".fasta"), "w") as f:
+                        skbio.io.write(genome, format="fasta", into=f)
+                    ids.add(fn)
+                else:
+                    duplicate_ids.add(fn)
+                    if error_on_duplicates:
+                        raise ValueError(msg.format(", ".join(duplicate_ids)))
+
+    else:
+        for genome in genomes:
+            for fp in genome.path.iterdir():
+                fn = os.path.basename(fp)
+                if fn not in ids:
+                    shutil.copyfile(
+                        fp,
+                        os.path.join(genomes_dir.path, fn),
+                    )
+                    ids.add(fn)
+                else:
+                    duplicate_ids.add(fn)
+                    if error_on_duplicates:
+                        raise ValueError(msg.format(", ".join(duplicate_ids)))
+
+    if duplicate_ids:
+        warn(
+            msg.format(", ".join(sorted(duplicate_ids)))
+            + " The latest occurrence will overwrite all previous "
+            "occurrences for each corresponding ID."
+        )
+
+    return genomes_dir
diff --git a/q2_assembly/plugin_setup.py b/q2_assembly/plugin_setup.py
@@ -418,6 +418,26 @@
     description="Not to be called directly. Used by map_reads.",
 )
 
+plugin.methods.register_function(
+    function=q2_assembly.helpers.collate_genomes,
+    inputs={"genomes": List[FeatureData[Sequence]] | List[GenomeData[DNASequence]]},
+    parameters={"on_duplicates": Str % Choices(["error", "warn"])},
+    outputs={"collated_genomes": GenomeData[DNASequence]},
+    input_descriptions={"genomes": "A  list of genomes to be collated."},
+    parameter_descriptions={
+        "on_duplicates": "Preferred behaviour when duplicated genome IDs "
+        'are encountered: "warn" displays a warning and '
+        "continues with the combination of the genomes "
+        'while "error" raises an error and aborts further '
+        "execution."
+    },
+    output_descriptions={"collated_genomes": "The converted genomes."},
+    name="Convert a list of FeatureData[Sequence] or a list of GenomeData[DNASequence] "
+    "to GenomeData[DNASequence].",
+    description="This method converts a list of FeatureData[Sequence] or a list of "
+    "GenomeData[DNASequence] to a GenomeData[DNASequence] artifact.",
+)
+
 plugin.register_semantic_types(QUASTResults)
 plugin.register_semantic_type_to_format(
     QUASTResults, artifact_format=QUASTResultsDirectoryFormat
diff --git a/q2_assembly/tests/data/dna-fasta-format/dna-sequences1.fasta b/q2_assembly/tests/data/dna-fasta-format/dna-sequences1.fasta
@@ -0,0 +1,4 @@
+>ref1 d_Bacteria_1
+ACGTACGT
+>ref2 d_Bacteria_2
+CGTCGTCC
diff --git a/q2_assembly/tests/data/dna-fasta-format/dna-sequences2.fasta b/q2_assembly/tests/data/dna-fasta-format/dna-sequences2.fasta
@@ -0,0 +1,4 @@
+>ref5 d_Bacteria_3
+ACGTACGT
+>ref6 d_Bacteria_4
+CGTCGTCC
diff --git a/q2_assembly/tests/data/genomes-dir-format1/ref1.fasta b/q2_assembly/tests/data/genomes-dir-format1/ref1.fasta
@@ -0,0 +1,2 @@
+>ref1
+ACGTTACGT
diff --git a/q2_assembly/tests/data/genomes-dir-format1/ref2.fasta b/q2_assembly/tests/data/genomes-dir-format1/ref2.fasta
@@ -0,0 +1,2 @@
+>ref2
+ACGGGTACT
diff --git a/q2_assembly/tests/data/genomes-dir-format2/ref3.fasta b/q2_assembly/tests/data/genomes-dir-format2/ref3.fasta
@@ -0,0 +1,2 @@
+>ref3
+ACGTTACGT
diff --git a/q2_assembly/tests/test_helpers.py b/q2_assembly/tests/test_helpers.py
@@ -13,15 +13,19 @@
 import tempfile
 import unittest
 import uuid
+import warnings
 from unittest.mock import ANY, call, patch
 
 import shortuuid
 import skbio
 from parameterized import parameterized
+from q2_types.feature_data import DNAFASTAFormat
+from q2_types.genome_data import GenomeSequencesDirectoryFormat
 from q2_types.per_sample_sequences import ContigSequencesDirFmt
 from qiime2.plugin.testing import TestPluginBase
+from qiime2.plugins import assembly
 
-from q2_assembly.helpers.helpers import rename_contigs
+from q2_assembly.helpers.helpers import collate_genomes, rename_contigs
 
 
 class TestUtils(TestPluginBase):
@@ -57,6 +61,141 @@ def test_is_valid_shortuuid(self):
         self.assertTrue(self.is_valid_shortuuid(true_shortuuid))
         self.assertFalse(self.is_valid_shortuuid(false_shortuuid))
 
+    @parameterized.expand(["single", "multiple"])
+    def test_collate_genomes_dnafastaformat(self, input):
+        genomes1 = DNAFASTAFormat(
+            self.get_data_path("dna-fasta-format/dna-sequences1.fasta"), "r"
+        )
+        genomes2 = DNAFASTAFormat(
+            self.get_data_path("dna-fasta-format/dna-sequences2.fasta"), "r"
+        )
+        if input == "single":
+            genomes = [genomes1]
+            content = {
+                "ref1": {"description": "d_Bacteria_1", "sequence": "ACGTACGT"},
+                "ref2": {"description": "d_Bacteria_2", "sequence": "CGTCGTCC"},
+            }
+            exp_files = ["ref1.fasta", "ref2.fasta"]
+        else:
+            genomes = [genomes1, genomes2]
+            content = {
+                "ref1": {"description": "d_Bacteria_1", "sequence": "ACGTACGT"},
+                "ref2": {"description": "d_Bacteria_2", "sequence": "CGTCGTCC"},
+                "ref5": {"description": "d_Bacteria_3", "sequence": "ACGTACGT"},
+                "ref6": {"description": "d_Bacteria_4", "sequence": "CGTCGTCC"},
+            }
+            exp_files = ["ref1.fasta", "ref2.fasta", "ref5.fasta", "ref6.fasta"]
+
+        collated_genomes = collate_genomes(genomes=genomes)
+        actual_files = sorted(os.listdir(collated_genomes.path))
+        self.assertEqual(actual_files, exp_files)
+
+        for fn in actual_files:
+            fp = os.path.join(collated_genomes.path, fn)
+            with open(fp, "r") as fasta_file:
+                for seq in skbio.io.read(fasta_file, "fasta"):
+                    actual_id = seq.metadata["id"]
+                    actual_description = seq.metadata["description"]
+                    actual_sequence = str(seq)
+                    expected_id = fn.split(".")[0]
+                    expected_desc = content[expected_id]["description"]
+                    expected_sequence = content[expected_id]["sequence"]
+
+                    self.assertEquals(actual_id, expected_id)
+                    self.assertEqual(actual_description, expected_desc)
+                    self.assertEqual(actual_sequence, expected_sequence)
+
+    def test_collate_genomes_genome_dir_multiple(self):
+        genomes1 = GenomeSequencesDirectoryFormat(
+            self.get_data_path("genomes-dir-format1"), "r"
+        )
+        genomes2 = GenomeSequencesDirectoryFormat(
+            self.get_data_path("genomes-dir-format2"), "r"
+        )
+        genomes = [genomes1, genomes2]
+        collated_genomes = collate_genomes(genomes=genomes)
+        exp_files = ["ref1.fasta", "ref2.fasta", "ref3.fasta"]
+        actual_files = sorted(os.listdir(collated_genomes.path))
+        self.assertEqual(exp_files, actual_files)
+
+    def test_collate_genomes_mix(self):
+        # should throw TypeError
+        genomes1 = DNAFASTAFormat(
+            self.get_data_path("dna-fasta-format/dna-sequences1.fasta"), "r"
+        )
+        genomes2 = GenomeSequencesDirectoryFormat(
+            self.get_data_path("genomes-dir-format2"), "r"
+        )
+        genomes = [genomes2, genomes1]
+        with self.assertRaises(TypeError):
+            assembly.methods.collate_genomes(genomes=genomes)
+
+    @parameterized.expand(["GenomeData", "DNAFASTAFormat"])
+    def test_collate_genomes_dnafastaformat_multiple_duplicates_warn(self, dir_fmt):
+        duplicate_ids = (
+            ["ref1.fasta", "ref2.fasta"]
+            if dir_fmt == "GenomeData"
+            else ["ref1", "ref2"]
+        )
+        warn_msg = (
+            "Duplicate sequence files were found for the following IDs: {}. "
+            "The latest occurrence will overwrite all previous occurrences "
+            "for each corresponding ID."
+        ).format(", ".join(duplicate_ids))
+        if dir_fmt == "GenomeData":
+            genomes1 = GenomeSequencesDirectoryFormat(
+                self.get_data_path("genomes-dir-format1"), "r"
+            )
+        else:
+            genomes1 = DNAFASTAFormat(
+                self.get_data_path("dna-fasta-format/dna-sequences1.fasta"), "r"
+            )
+        with warnings.catch_warnings(record=True) as w:
+            collated_genomes = collate_genomes(genomes=[genomes1, genomes1])
+            exp_files = ["ref1.fasta", "ref2.fasta"]
+            actual_files = sorted(os.listdir(collated_genomes.path))
+            self.assertEqual(actual_files, exp_files)
+            self.assertEqual(warn_msg, str(w[0].message))
+
+            if dir_fmt == "DNAFASTAFormat":
+                content = {
+                    "ref1": {"description": "d_Bacteria_1", "sequence": "ACGTACGT"},
+                    "ref2": {"description": "d_Bacteria_2", "sequence": "CGTCGTCC"},
+                }
+
+                for fn in actual_files:
+                    fp = os.path.join(collated_genomes.path, fn)
+                    with open(fp, "r") as fasta_file:
+                        for seq in skbio.io.read(fasta_file, "fasta"):
+                            actual_id = seq.metadata["id"]
+                            actual_description = seq.metadata["description"]
+                            actual_sequence = str(seq)
+                            expected_id = fn.split(".")[0]
+                            expected_desc = content[expected_id]["description"]
+                            expected_sequence = content[expected_id]["sequence"]
+
+                            self.assertEquals(actual_id, expected_id)
+                            self.assertEqual(actual_description, expected_desc)
+                            self.assertEqual(actual_sequence, expected_sequence)
+
+    @parameterized.expand(["GenomeData", "DNAFASTAFormat"])
+    def test_collate_genomes_duplicates_error(self, dir_fmt):
+        duplicate_ids = ["ref3.fasta"] if dir_fmt == "GenomeData" else ["ref1"]
+        error_msg = (
+            "Duplicate sequence files were found for the "
+            "following IDs: %s." % ", ".join(duplicate_ids)
+        )
+        if dir_fmt == "GenomeData":
+            genomes1 = GenomeSequencesDirectoryFormat(
+                self.get_data_path("genomes-dir-format2"), "r"
+            )
+        else:
+            genomes1 = DNAFASTAFormat(
+                self.get_data_path("dna-fasta-format/dna-sequences1.fasta"), "r"
+            )
+        with self.assertRaisesRegex(ValueError, error_msg):
+            _ = collate_genomes(genomes=[genomes1, genomes1], on_duplicates="error")
+
     @parameterized.expand(
         [
             ("uuid4", UUID4_REGEX),
diff --git a/setup.py b/setup.py
@@ -43,6 +43,9 @@
             "data/zip_test_data/expected/*/*/*/*",
             "data/formatted-reads/single-end/*",
             "data/formatted-reads/paired-end/*",
+            "data/dna-fasta-format/*",
+            "data/genomes-dir-format1/*",
+            "data/genomes-dir-format2/*",
         ],
         "q2_assembly.bowtie2.tests": [
             "data/*",
