From e8203ad5948f5428a89086690d5d8e954ef0d9a1 Mon Sep 17 00:00:00 2001
From: VinzentRisch <risch.vinzent@gmail.com>
Date: Wed, 20 Mar 2024 14:11:49 +0100
Subject: [PATCH] added type map for collate of allele annotations

---
 q2_amr/plugin_setup.py | 30 ++++++++++++++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

diff --git a/q2_amr/plugin_setup.py b/q2_amr/plugin_setup.py
index d507045..df15724 100644
--- a/q2_amr/plugin_setup.py
+++ b/q2_amr/plugin_setup.py
@@ -246,11 +246,37 @@
     "and collates them into a single artifact.",
 )
 
+T_allele_annotation_collate_in, T_allele_annotation_collate_out = TypeMap(
+    {
+        List[
+            SampleData[CARDAlleleAnnotation % Properties("kma", "bowtie2", "bwa")]
+        ]: SampleData[CARDAlleleAnnotation % Properties("kma", "bowtie2", "bwa")],
+        List[
+            SampleData[CARDAlleleAnnotation % Properties("kma", "bowtie2")]
+        ]: SampleData[CARDAlleleAnnotation % Properties("kma", "bowtie2")],
+        List[SampleData[CARDAlleleAnnotation % Properties("kma", "bwa")]]: SampleData[
+            CARDAlleleAnnotation % Properties("kma", "bwa")
+        ],
+        List[
+            SampleData[CARDAlleleAnnotation % Properties("bowtie2", "bwa")]
+        ]: SampleData[CARDAlleleAnnotation % Properties("bowtie2", "bwa")],
+        List[SampleData[CARDAlleleAnnotation % Properties("kma")]]: SampleData[
+            CARDAlleleAnnotation % Properties("kma")
+        ],
+        List[SampleData[CARDAlleleAnnotation % Properties("bowtie2")]]: SampleData[
+            CARDAlleleAnnotation % Properties("bowtie2")
+        ],
+        List[SampleData[CARDAlleleAnnotation % Properties("bwa")]]: SampleData[
+            CARDAlleleAnnotation % Properties("bwa")
+        ],
+    }
+)
+
 plugin.methods.register_function(
     function=collate_reads_allele_annotations,
-    inputs={"annotations": List[SampleData[CARDAlleleAnnotation]]},
+    inputs={"annotations": T_allele_annotation_collate_in},
     parameters={},
-    outputs={"collated_annotations": SampleData[CARDAlleleAnnotation]},
+    outputs={"collated_annotations": T_allele_annotation_collate_out},
     input_descriptions={
         "annotations": "A collection of annotations from reads at "
         "allele level to be collated."