added sampledata and feature data dir fmts

q2_amr/amrfinderplus/types/_format.py

@@ -5,8 +5,11 @@
 #
 # The full license is in the file LICENSE, distributed with this software.
 # ----------------------------------------------------------------------------
+import os
+
 import pandas as pd
 from q2_types.feature_data import MixedCaseDNAFASTAFormat, ProteinFASTAFormat
+from q2_types.per_sample_sequences._format import MultiDirValidationMixin
 from qiime2.core.exceptions import ValidationError
 from qiime2.plugin import model
 
@@ -389,3 +392,29 @@ def _validate(self, n_records=None):
 
     def _validate_(self, level):
         self._validate()
+
+
+class ARMFinderPlusAnnotationsDirFmt(MultiDirValidationMixin, model.DirectoryFormat):
+    tsv = model.FileCollection(
+        r".+amr_annotation.tsv$", format=ARMFinderPlusAnnotationFormat
+    )
+
+    @tsv.set_path_maker
+    def json_path_maker(self, sample_id, mag_id):
+        return f"{sample_id}/{mag_id}/amr_annotation.tsv$"
+
+    def sample_dict(self):
+        sample_dict = {}
+        for sample in self.path.iterdir():
+            mag_dict = {}
+            for mag in sample.iterdir():
+                mag_dict[mag.name] = [
+                    os.path.join(mag, "amr_annotation.tsv"),
+                ]
+            sample_dict[sample.name] = mag_dict
+        return sample_dict
+
+
+ARMFinderPlusAnnotationDirFmt = model.SingleFileDirectoryFormat(
+    "ARMFinderPlusAnnotationDirFmt", "amr_annotation.tsv", ARMFinderPlusAnnotationFormat
+)

q2_amr/amrfinderplus/types/_type.py

@@ -5,7 +5,14 @@
 #
 # The full license is in the file LICENSE, distributed with this software.
 # ----------------------------------------------------------------------------
+from q2_types.feature_data import FeatureData
+from q2_types.sample_data import SampleData
 from qiime2.core.type import SemanticType
 
 AMRFinderPlusDatabase = SemanticType("AMRFinderPlusDatabase")
-ARMFinderPlusAnnotation = SemanticType("ARMFinderPlusAnnotation")
+ARMFinderPlusAnnotations = SemanticType(
+    "ARMFinderPlusAnnotations", variant_of=SampleData.field["type"]
+)
+ARMFinderPlusAnnotation = SemanticType(
+    "ARMFinderPlusAnnotation", variant_of=FeatureData.field["type"]
+)

q2_amr/amrfinderplus/types/tests/test_types_formats_transformers.py

@@ -5,12 +5,16 @@
 #
 # The full license is in the file LICENSE, distributed with this software.
 # ----------------------------------------------------------------------------
+import os
+
 from qiime2.core.exceptions import ValidationError
 from qiime2.plugin.testing import TestPluginBase
 
 from q2_amr.amrfinderplus.types._format import (
     AMRFinderPlusDatabaseDirectoryFormat,
+    ARMFinderPlusAnnotationDirFmt,
     ARMFinderPlusAnnotationFormat,
+    ARMFinderPlusAnnotationsDirFmt,
 )
 
 
@@ -24,18 +28,25 @@ def test_amrfinderplus_database_directory_format_validate_positive(self):
         format.validate()
 
     def test_amrfinderplus_annotation_format_validate_positive(self):
-        filepath = self.get_data_path("annotation/amr_annotation.tsv")
+        filepath = self.get_data_path(
+            "annotation/no_coordinates/aa447c99-ecd9-4c4a-a53b-4df6999815dd"
+            "/amr_annotation.tsv"
+        )
+
         format = ARMFinderPlusAnnotationFormat(filepath, mode="r")
         format.validate()
 
     def test_amrfinderplus_annotation_format_validate_positive_coordinates(self):
-        filepath = self.get_data_path("annotation/amr_annotation_coordiantes.tsv")
+        filepath = self.get_data_path(
+            "annotation/coordinates/e026af61-d911-4de3-a957-7e8bf837f30d"
+            "/amr_annotation.tsv"
+        )
         format = ARMFinderPlusAnnotationFormat(filepath, mode="r")
         format.validate()
 
     def test_amrfinderplus_annotation_format_validation_error(self):
         with self.assertRaises(ValidationError) as context:
-            path = self.get_data_path("annotation/amr_annotation_wrong.tsv")
+            path = self.get_data_path("annotation_wrong/amr_annotation.tsv")
             format = ARMFinderPlusAnnotationFormat(path, mode="r")
             format.validate()
 
@@ -73,3 +84,40 @@ def test_amrfinderplus_annotation_format_validation_error(self):
             )
 
             self.assertEqual(str(context.exception), expected_message)
+
+    def test_amrfinderplus_annotations_directory_format_sample_dict(self):
+        dirpath = self.get_data_path("annotation")
+        annotations = ARMFinderPlusAnnotationsDirFmt(dirpath, mode="r")
+
+        obs = annotations.sample_dict()
+
+        exp = {
+            "coordinates": {
+                "e026af61-d911-4de3-a957-7e8bf837f30d": [
+                    os.path.join(
+                        annotations.path,
+                        "coordinates",
+                        "e026af61-d911-4de3-a957-7e8bf837f30d",
+                        "amr_annotation.tsv",
+                    ),
+                ]
+            },
+            "no_coordinates": {
+                "aa447c99-ecd9-4c4a-a53b-4df6999815dd": [
+                    os.path.join(
+                        annotations.path,
+                        "no_coordinates",
+                        "aa447c99-ecd9-4c4a-a53b-4df6999815dd",
+                        "amr_annotation.tsv",
+                    ),
+                ],
+            },
+        }
+        self.assertEqual(obs, exp)
+
+    def test_amrfinderplus_annotation_directory_format(self):
+        dirpath = self.get_data_path(
+            "annotation/coordinates/e026af61-d911-4de3-a957-7e8bf837f30d"
+        )
+        annotations = ARMFinderPlusAnnotationDirFmt(dirpath, mode="r")
+        assert isinstance(annotations, ARMFinderPlusAnnotationDirFmt)

q2_amr/plugin_setup.py

@@ -7,6 +7,7 @@
 # ----------------------------------------------------------------------------
 import importlib
 
+from q2_types.feature_data import FeatureData
 from q2_types.feature_table import FeatureTable, Frequency
 from q2_types.per_sample_sequences import (
     MAGs,
@@ -30,11 +31,17 @@
 from q2_amr import __version__
 from q2_amr.amrfinderplus.types._format import (
     AMRFinderPlusDatabaseDirectoryFormat,
+    ARMFinderPlusAnnotationDirFmt,
     ARMFinderPlusAnnotationFormat,
+    ARMFinderPlusAnnotationsDirFmt,
     BinaryFormat,
     TextFormat,
 )
-from q2_amr.amrfinderplus.types._type import AMRFinderPlusDatabase
+from q2_amr.amrfinderplus.types._type import (
+    AMRFinderPlusDatabase,
+    ARMFinderPlusAnnotation,
+    ARMFinderPlusAnnotations,
+)
 from q2_amr.card.database import fetch_card_db
 from q2_amr.card.heatmap import heatmap
 from q2_amr.card.kmer import (
@@ -1119,6 +1126,14 @@
     AMRFinderPlusDatabase,
     artifact_format=AMRFinderPlusDatabaseDirectoryFormat,
 )
+plugin.register_semantic_type_to_format(
+    SampleData[ARMFinderPlusAnnotations],
+    artifact_format=ARMFinderPlusAnnotationsDirFmt,
+)
+plugin.register_semantic_type_to_format(
+    FeatureData[ARMFinderPlusAnnotation],
+    artifact_format=ARMFinderPlusAnnotationDirFmt,
+)
 plugin.register_formats(
     CARDKmerDatabaseDirectoryFormat,
     CARDKmerJSONFormat,