update plugin_setup

bokulich-lab · May 10, 2024 · 61d128d · 61d128d
1 parent 7ea1233
commit 61d128d
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 2 deletions.
diff --git a/q2_moshpit/busco/tests/test_fetch_busco.py b/q2_moshpit/busco/tests/test_fetch_busco.py
@@ -35,4 +35,4 @@ def test_fetch_busco_db_all(self, subp_run):
 
         # Check that command was called in the expected way
         cmd = ["busco", "--download", "all"]
-        subp_run.assert_called_once_with(cmd, check=True, cwd=str(busco_db))
+        subp_run.assert_called_once_with(cmd, check=True, cwd=str(busco_db))
diff --git a/q2_moshpit/plugin_setup.py b/q2_moshpit/plugin_setup.py
@@ -38,7 +38,8 @@
 from q2_types.kraken2._type import BrackenDB
 from q2_types.per_sample_sequences._type import AlignmentMap
 from q2_types.reference_db import (
-    ReferenceDB, Diamond, Eggnog, NCBITaxonomy, EggnogProteinSequences
+    ReferenceDB, Diamond, Eggnog, NCBITaxonomy, EggnogProteinSequences,
+    BuscoDB
 )
 
 citations = Citations.load('citations.bib', package='q2_moshpit')
@@ -1200,3 +1201,65 @@
     artifact_format=BUSCOResultsDirectoryFormat)
 plugin.register_formats(BUSCOResultsFormat, BUSCOResultsDirectoryFormat)
 importlib.import_module('q2_moshpit.busco.types._transformer')
+
+p_virus, p_prok, p_euk, o_busco_db = TypeMap({
+    (
+        Bool % Choices(True),
+        Bool % Choices(True),
+        Bool % Choices(True)
+    ): ReferenceDB[BuscoDB % Properties(['virus', 'prokaryota', 'eukaryota'])],
+    (
+        Bool % Choices(False),
+        Bool % Choices(True),
+        Bool % Choices(True)
+    ): ReferenceDB[BuscoDB % Properties(['prokaryota', 'eukaryota'])],
+    (
+        Bool % Choices(True),
+        Bool % Choices(False),
+        Bool % Choices(True)
+    ): ReferenceDB[BuscoDB % Properties(['virus', 'eukaryota'])],
+    (
+        Bool % Choices(True),
+        Bool % Choices(True),
+        Bool % Choices(False)
+    ): ReferenceDB[BuscoDB % Properties(['virus', 'prokaryota'])],
+    (
+        Bool % Choices(True),
+        Bool % Choices(False),
+        Bool % Choices(False)
+    ): ReferenceDB[BuscoDB % Properties('virus')],
+    (
+        Bool % Choices(False),
+        Bool % Choices(True),
+        Bool % Choices(False)
+    ): ReferenceDB[BuscoDB % Properties('prokaryota')],
+    (
+        Bool % Choices(False),
+        Bool % Choices(False),
+        Bool % Choices(True)
+    ): ReferenceDB[BuscoDB % Properties('eukaryota')],
+})
+
+plugin.methods.register_function(
+    function=q2_moshpit.busco.fetch_busco_db,
+    inputs={},
+    outputs=[('busco_db', o_busco_db)],
+    output_descriptions={
+        'busco_db': "BUSCO database for the specified lineages"
+    },
+    parameters={
+        "virus": p_virus,
+        "prok": p_prok,
+        "euk": p_euk,
+    },
+    parameter_descriptions={
+        "virus": "Download the virus dataset",
+        "prok": "Download the prokaryote dataset",
+        "euk": "Download the eukaryote dataset",
+    },
+    name="Download BUSCO database.",
+    description="Downloads BUSCO database for the specified lineage. "
+                "Output can be used to run BUSCO with the 'evaluate-busco' "
+                "action.",
+    citations=[citations["manni_busco_2021"]],
+)