ENH: add collate-annotations action

q2_moshpit/partition/__init__.py

@@ -5,6 +5,7 @@
 #
 # The full license is in the file LICENSE, distributed with this software.
 # ----------------------------------------------------------------------------
+from .annotations import collate_annotations
 from .mags import (
     collate_feature_data_mags, collate_sample_data_mags,
     partition_feature_data_mags, partition_sample_data_mags
@@ -14,5 +15,5 @@
 __all__ = [
     "collate_feature_data_mags", "collate_sample_data_mags",
     "partition_feature_data_mags", "partition_sample_data_mags",
-    "collate_orthologs"
+    "collate_orthologs", "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

@@ -1,5 +1,5 @@
 # ----------------------------------------------------------------------------
-# Copyright (c) 2023, QIIME 2 development team.
+# Copyright (c) 2023, QIIME 2 development team.n
 #
 # Distributed under the terms of the Modified BSD License.
 #

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

@@ -835,6 +835,22 @@
     description="Collate a List of SampleData[BLAST6] into one"
 )
 
+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 ortholog annotations."
+    },
+    output_descriptions={
+        'collated_annotations': "Collated ortholog annotations."
+    },
+    name='Collate ortholog annotations.',
+    description="Takes a collection of FeatureData[NOG]'s "
+                "and collates them into a single artifact.",
+)
+
 busco_params = {
     "mode": Str % Choices(["genome"]),
     "lineage_dataset": Str,