ENH: add collate-annotations action

q2_moshpit/__init__.py

@@ -15,6 +15,7 @@
 from . import prodigal
 from . import eggnog
 from . import busco
+from . import partition
 
 
 from ._version import get_versions
@@ -24,5 +25,5 @@
 __all__ = [
     'metabat2', 'bracken', 'kraken_class', 'kraken_db',
     'kaiju_class', 'kaiju_db', 'dereplicate_mags', 'eggnog',
-    'busco', 'prodigal', 'helpers'
+    'busco', 'prodigal', 'helpers', 'partition'
 ]

q2_moshpit/partition/__init__.py

@@ -0,0 +1,10 @@
+# ----------------------------------------------------------------------------
+# Copyright (c) 2022-2023, QIIME 2 development team.
+#
+# Distributed under the terms of the Modified BSD License.
+#
+# The full license is in the file LICENSE, distributed with this software.
+# ----------------------------------------------------------------------------
+from .annotations import collate_annotations
+
+__all__ = ["collate_annotations"]

q2_moshpit/partition/annotations.py

@@ -0,0 +1,23 @@
+# ----------------------------------------------------------------------------
+# Copyright (c) 2022-2023, QIIME 2 development team.
+#
+# Distributed under the terms of the Modified BSD License.
+#
+# The full license is in the file LICENSE, distributed with this software.
+# ----------------------------------------------------------------------------
+from q2_types.feature_data_mag import OrthologAnnotationDirFmt
+from qiime2.util import duplicate
+
+
+def collate_annotations(
+    ortholog_annotations: OrthologAnnotationDirFmt
+) -> OrthologAnnotationDirFmt:
+    # Init output
+    collated_annotations = OrthologAnnotationDirFmt()
+
+    # Copy annotations into output
+    for anno in ortholog_annotations:
+        for fp in anno.path.iterdir():
+            duplicate(fp, collated_annotations.path / fp.name)
+
+    return collated_annotations

q2_moshpit/partition/tests/__init__.py

@@ -0,0 +1,7 @@
+# ----------------------------------------------------------------------------
+# Copyright (c) 2022-2023, QIIME 2 development team.
+#
+# Distributed under the terms of the Modified BSD License.
+#
+# The full license is in the file LICENSE, distributed with this software.
+# ----------------------------------------------------------------------------

q2_moshpit/partition/tests/test_annotations.py

@@ -0,0 +1,33 @@
+# ----------------------------------------------------------------------------
+# Copyright (c) 2022-2023, QIIME 2 development team.
+#
+# Distributed under the terms of the Modified BSD License.
+#
+# The full license is in the file LICENSE, distributed with this software.
+# ----------------------------------------------------------------------------
+import filecmp
+from qiime2.plugin.testing import TestPluginBase
+from q2_types.feature_data_mag import OrthologAnnotationDirFmt
+from ..annotations import collate_annotations
+
+
+class TestCollateAnnotations(TestPluginBase):
+    package = 'q2_moshpit.partition.tests'
+
+    def test_collate_annotations(self):
+        p = self.get_data_path("annotations")
+        annotations = [
+          OrthologAnnotationDirFmt(f"{p}/{letter}", mode="r")
+          for letter in ["a", "b", "c"]
+        ]
+        collated_annotations = collate_annotations(annotations)
+
+        # assert that all files are there
+        compare = filecmp.dircmp(
+            collated_annotations.path,
+            self.get_data_path("annotations/collated")
+        )
+        self.assertListEqual(
+            compare.common,
+            [f"{letter}.annotations" for letter in ["a", "b", "c"]]
+        )

q2_moshpit/plugin_setup.py

@@ -710,6 +710,23 @@
     citations=[citations["huerta_cepas_eggnog_2019"]]
 )
 
+plugin.methods.register_function(
+    function=q2_moshpit.partition.collate_annotations,
+    inputs={'ortholog_annotations': List[FeatureData[NOG]]},
+    parameters={},
+    outputs=[('collated_annotations', FeatureData[NOG])],
+    input_descriptions={
+        'ortholog_annotations': "Collection of orthologs annotations"
+    },
+    output_descriptions={
+        'collated_annotations': "Collated orthologs annotations"
+    },
+    name='Collate orthologs annotations',
+    description="Collate orthologs annotations from the eggnog-annotate "
+                "action. Intended when parallel runs of eggnog-diamond-search "
+                "and eggnog-annotate were used.",
+)
+
 busco_params = {
     "mode": Str % Choices(["genome"]),
     "lineage_dataset": Str,