ENH: Add aligner property to annotate-reads-card outputs (#20)

bokulich-lab · Jan 16, 2024 · c0ca96d · c0ca96d
1 parent d7f131c
commit c0ca96d
Showing 1 changed file with 24 additions and 4 deletions.
diff --git 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,21 +118,41 @@
     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={
         "reads": SampleData[PairedEndSequencesWithQuality | SequencesWithQuality],
         "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]),
     ],