ENH: Added sample_dict to CARDAlleleAnnotationDirectoryFormat , CARDGeneAnnotationDirectoryFormat and CARDAnnotationDirectoryFormat

q2_amr/types/_format.py

@@ -6,6 +6,7 @@
 # The full license is in the file LICENSE, distributed with this software.
 # ----------------------------------------------------------------------------
 import json
+import os
 import re
 from copy import copy
 
@@ -277,6 +278,20 @@ def json_path_maker(self, sample_id, bin_id):
     def txt_path_maker(self, sample_id, bin_id):
         return f"{sample_id}/{bin_id}/amr_annotation.txt"
 
+    def sample_dict(self):
+        sample_dict = {}
+        for sample in self.path.iterdir():
+            for mag in sample.iterdir():
+                files = [
+                    os.path.join(mag, file)
+                    for file in [
+                        "amr_annotation.json",
+                        "amr_annotation.txt",
+                    ]
+                ]
+                sample_dict[sample.name] = {mag.name: files}
+        return sample_dict
+
 
 class CARDAlleleAnnotationFormat(model.TextFileFormat):
     def _validate(self, n_records=None):
@@ -421,6 +436,20 @@ def stats_path_maker(self, sample_id):
     def bam_path_maker(self, sample_id):
         return "%s/sorted.length_100.bam" % sample_id
 
+    def sample_dict(self):
+        sample_dict = {}
+        for sample in self.path.iterdir():
+            files = [
+                os.path.join(sample, file)
+                for file in [
+                    "allele_mapping_data.txt",
+                    "overall_mapping_stats.txt",
+                    "sorted.length_100.bam",
+                ]
+            ]
+            sample_dict[sample.name] = files
+        return sample_dict
+
 
 class CARDGeneAnnotationDirectoryFormat(MultiDirValidationMixin, model.DirectoryFormat):
     gene = model.FileCollection(
@@ -431,6 +460,13 @@ class CARDGeneAnnotationDirectoryFormat(MultiDirValidationMixin, model.Directory
     def gene_path_maker(self, sample_id):
         return "%s/gene_mapping_data.txt" % sample_id
 
+    def sample_dict(self):
+        sample_dict = {}
+        for sample in self.path.iterdir():
+            file = list(os.path.join(sample, "gene_mapping_data.txt"))
+            sample_dict[sample.name] = file
+        return sample_dict
+
 
 class CARDMAGsKmerAnalysisFormat(model.TextFileFormat):
     def _validate(self, n_records=None):

q2_amr/types/tests/test_types_formats_transformers.py

@@ -315,28 +315,78 @@ def test_CARDAnnotationDirectoryFormat_to_qiime2_Metadata_transformer(self):
         metadata_obt = transformer(annotation)
         self.assertIsInstance(metadata_obt, qiime2.Metadata)
 
+    def test_card_annotation_directory_format_sample_dict(self):
+        dirpath = self.get_data_path("annotate_mags_output")
+        annotations = CARDAnnotationDirectoryFormat(dirpath, mode="r")
+
+        obs = annotations.sample_dict()
+        exp = {
+            "sample1": {
+                "bin1": [
+                    os.path.join(dirpath, "sample1", "bin1", "amr_annotation.json"),
+                    os.path.join(dirpath, "sample1", "bin1", "amr_annotation.txt"),
+                ]
+            },
+            "sample2": {
+                "bin1": [
+                    os.path.join(dirpath, "sample2", "bin1", "amr_annotation.json"),
+                    os.path.join(dirpath, "sample2", "bin1", "amr_annotation.txt"),
+                ]
+            },
+        }
+        self.assertEqual(obs, exp)
+
 
 class TestCARDReadsAnnotationTypesAndFormats(AMRTypesTestPluginBase):
     def test_CARDGeneAnnotationDirectoryFormat_to_qiime2_Metadata_transformer(self):
         transformer = self.get_transformer(
             CARDGeneAnnotationDirectoryFormat, qiime2.Metadata
         )
         annotation = CARDGeneAnnotationDirectoryFormat(
-            self.get_data_path("annotate_reads_output"), "r"
+            self.get_data_path("card_gene_annotation"), "r"
         )
         metadata_obt = transformer(annotation)
         self.assertIsInstance(metadata_obt, qiime2.Metadata)
 
     def test_CARDAlleleAnnotationDirectoryFormat_to_qiime2_Metadata_transformer(self):
         transformer = self.get_transformer(
-            CARDGeneAnnotationDirectoryFormat, qiime2.Metadata
+            CARDAlleleAnnotationDirectoryFormat, qiime2.Metadata
         )
         annotation = CARDAlleleAnnotationDirectoryFormat(
-            self.get_data_path("annotate_reads_output"), "r"
+            self.get_data_path("card_allele_annotation"), "r"
         )
         metadata_obt = transformer(annotation)
         self.assertIsInstance(metadata_obt, qiime2.Metadata)
 
+    def test_card_allele_annotation_directory_format_sample_dict(self):
+        dirpath = self.get_data_path("card_allele_annotation")
+        annotations = CARDAlleleAnnotationDirectoryFormat(dirpath, mode="r")
+
+        obs = annotations.sample_dict()
+        exp = {
+            sample: [
+                os.path.join(dirpath, sample, file)
+                for file in [
+                    "allele_mapping_data.txt",
+                    "overall_mapping_stats.txt",
+                    "sorted.length_100.bam",
+                ]
+            ]
+            for sample in ["sample1", "sample2"]
+        }
+        self.assertEqual(obs, exp)
+
+    def test_card_gene_annotation_directory_format_sample_dict(self):
+        dirpath = self.get_data_path("card_gene_annotation")
+        annotations = CARDGeneAnnotationDirectoryFormat(dirpath, mode="r")
+
+        obs = annotations.sample_dict()
+        exp = {
+            "sample1": list(os.path.join(dirpath, "sample1", "gene_mapping_data.txt")),
+            "sample2": list(os.path.join(dirpath, "sample2", "gene_mapping_data.txt")),
+        }
+        self.assertEqual(obs, exp)
+
 
 class TestKmerTypesAndFormats(AMRTypesTestPluginBase):
     def test_card_mags_kmer_analysis_validate_positive(self):

setup.py

@@ -31,7 +31,8 @@
         "q2_amr.types.tests": [
             "data/*",
             "data/annotate_mags_output/*/*/*",
-            "data/annotate_reads_output/*/*",
+            "data/card_allele_annotation/*/*",
+            "data/card_gene_annotation/*/*",
         ],
         "q2_amr.card.tests": ["data/*"],
     },