diff --git a/q2_amr/plugin_setup.py b/q2_amr/plugin_setup.py index 941404d..608d0b2 100644 --- a/q2_amr/plugin_setup.py +++ b/q2_amr/plugin_setup.py @@ -14,7 +14,7 @@ ) from q2_types.sample_data import SampleData from q2_types_genomics.per_sample_data import MAGs -from qiime2.core.type import Bool, Choices, Int, Range, Str +from qiime2.core.type import Bool, Choices, Int, Properties, Range, Str, TypeMap from qiime2.plugin import Citations, Plugin from q2_amr import __version__ @@ -118,6 +118,26 @@ citations=[citations["alcock_card_2023"]], ) +P_aligner, T_allele_annotation, T_gene_annotation = TypeMap( + { + Str + % Choices("kma"): ( + SampleData[CARDAlleleAnnotation % Properties("kma")], + SampleData[CARDGeneAnnotation % Properties("kma")], + ), + Str + % Choices("bowtie2"): ( + SampleData[CARDAlleleAnnotation % Properties("bowtie2")], + SampleData[CARDGeneAnnotation % Properties("bowtie2")], + ), + Str + % Choices("bwa"): ( + SampleData[CARDAlleleAnnotation % Properties("bwa")], + SampleData[CARDGeneAnnotation % Properties("bwa")], + ), + } +) + plugin.methods.register_function( function=annotate_reads_card, inputs={ @@ -125,14 +145,14 @@ "card_db": CARDDatabase, }, parameters={ - "aligner": Str % Choices(["kma", "bowtie2", "bwa"]), + "aligner": P_aligner, "threads": Int % Range(0, None, inclusive_start=False), "include_wildcard": Bool, "include_other_models": Bool, }, outputs=[ - ("amr_allele_annotation", SampleData[CARDAlleleAnnotation]), - ("amr_gene_annotation", SampleData[CARDGeneAnnotation]), + ("amr_allele_annotation", T_allele_annotation), + ("amr_gene_annotation", T_gene_annotation), ("allele_feature_table", FeatureTable[Frequency]), ("gene_feature_table", FeatureTable[Frequency]), ],