From cc386b350c2e7d17b3bd7874962b215e09994d00 Mon Sep 17 00:00:00 2001
From: VinzentRisch <risch.vinzent@gmail.com>
Date: Wed, 10 Jan 2024 17:13:33 +0100
Subject: [PATCH] MAINT: Removed overall-mapping-stats.txt from
 CARDGeneAnnotationDirectoryFormat

---
 q2_amr/card/reads.py            | 11 ++++++-----
 q2_amr/plugin_setup.py          |  6 ++----
 q2_amr/tests/card/test_reads.py |  6 ++++--
 q2_amr/types/_format.py         |  7 -------
 4 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/q2_amr/card/reads.py b/q2_amr/card/reads.py
index 5a97bca..ec8ead7 100644
--- a/q2_amr/card/reads.py
+++ b/q2_amr/card/reads.py
@@ -85,8 +85,11 @@ def annotate_reads_card(
             for map_type, des_dir in zip(
                 ["allele", "gene"], [samp_allele_dir, samp_gene_dir]
             ):
-                files = [f"{map_type}_mapping_data.txt", "overall_mapping_stats.txt"]
-
+                files = [f"{map_type}_mapping_data.txt"]
+                # mapping statistics only go to the allele directories
+                files.append(
+                    "overall_mapping_stats.txt"
+                ) if map_type == "allele" else None
                 for file in files:
                     shutil.copy(
                         os.path.join(samp_input_dir, "output." + file),
@@ -215,9 +218,7 @@ def extract_sample_stats(samp_dir: str):
 
 def visualize_annotation_stats(
     output_dir: str,
-    amr_reads_annotation: Union[
-        CARDGeneAnnotationDirectoryFormat, CARDAlleleAnnotationDirectoryFormat
-    ],
+    amr_reads_annotation: CARDAlleleAnnotationDirectoryFormat,
 ):
     directory = str(amr_reads_annotation)
     sample_stats = {}
diff --git a/q2_amr/plugin_setup.py b/q2_amr/plugin_setup.py
index 413e850..941404d 100644
--- a/q2_amr/plugin_setup.py
+++ b/q2_amr/plugin_setup.py
@@ -196,11 +196,9 @@
 
 plugin.visualizers.register_function(
     function=visualize_annotation_stats,
-    inputs={"amr_reads_annotation": CARDGeneAnnotation | CARDAlleleAnnotation},
+    inputs={"amr_reads_annotation": SampleData[CARDAlleleAnnotation]},
     parameters={},
-    input_descriptions={
-        "amr_reads_annotation": "AMR annotation mapped on alleles or genes."
-    },
+    input_descriptions={"amr_reads_annotation": "AMR annotations on the allele level."},
     parameter_descriptions={},
     name="Visualize mapping statistics.",
     description="Visualize mapping statistics of an annotate-reads-card output.",
diff --git a/q2_amr/tests/card/test_reads.py b/q2_amr/tests/card/test_reads.py
index bacca85..151b773 100644
--- a/q2_amr/tests/card/test_reads.py
+++ b/q2_amr/tests/card/test_reads.py
@@ -144,10 +144,12 @@ def annotate_reads_card_test_body(self, read_type):
 
             # Assert if the expected files are in every sample directory and in both
             # resulting CARD annotation objects
-
             for num in [0, 1]:
                 map_type = "allele" if num == 0 else "gene"
-                files = [f"{map_type}_mapping_data.txt", "overall_mapping_stats.txt"]
+                files = [f"{map_type}_mapping_data.txt"]
+                files.append(
+                    "overall_mapping_stats.txt"
+                ) if map_type == "allele" else None
                 for samp in ["sample1", "sample2"]:
                     for file in files:
                         self.assertTrue(
diff --git a/q2_amr/types/_format.py b/q2_amr/types/_format.py
index 2eb3d36..de8c44a 100644
--- a/q2_amr/types/_format.py
+++ b/q2_amr/types/_format.py
@@ -421,14 +421,7 @@ class CARDGeneAnnotationDirectoryFormat(MultiDirValidationMixin, model.Directory
     gene = model.FileCollection(
         r".+(gene_mapping_data.txt)$", format=CARDGeneAnnotationFormat
     )
-    stats = model.FileCollection(
-        r".+(overall_mapping_stats.txt)$", format=CARDAnnotationStatsFormat
-    )
 
     @gene.set_path_maker
     def gene_path_maker(self, sample_id):
         return "%s/gene_mapping_data.txt" % sample_id
-
-    @stats.set_path_maker
-    def stats_path_maker(self, sample_id):
-        return "%s/overall_mapping_stats.txt" % sample_id