Add a file_dict method to the BamDirFmt and MultiBamDirFmt (#387)

* Add a file_dict method to the BamDirFmt and MultiBamDirFmt

* Shorten long line

Author: fethalen

q2_types/per_sample_sequences/_formats.py

@@ -24,7 +24,7 @@
 from q2_types.bowtie2 import Bowtie2IndexDirFmt
 from q2_types.feature_data import DNAFASTAFormat
 from ._util import _parse_sequence_filename, _manifest_to_df
-from .._util import FastqGzFormat
+from .._util import FastqGzFormat, FileDictMixin
 from ._util import validate_paired_ends_equal_record_count
 
 
@@ -708,15 +708,20 @@ def _validate_(self, level):
 
 
 # borrowed from q2-phylogenomics
-class BAMDirFmt(model.DirectoryFormat):
+class BAMDirFmt(model.DirectoryFormat, FileDictMixin):
+    pathspec = r'.+\.bam$'
+    suffixes = ['.bam']
     bams = model.FileCollection(r'.+\.bam', format=BAMFormat)
 
     @bams.set_path_maker
     def bams_path_maker(self, sample_id):
         return '%s.bam' % sample_id
 
 
-class MultiBAMDirFmt(MultiDirValidationMixin, model.DirectoryFormat):
+class MultiBAMDirFmt(MultiDirValidationMixin, model.DirectoryFormat,
+                     FileDictMixin):
+    pathspec = r'.+\.bam$'
+    suffixes = ['.bam']
     bams = model.FileCollection(r'.+\/.+\.bam', format=BAMFormat)
 
     @bams.set_path_maker

q2_types/per_sample_sequences/tests/test_formats.py

@@ -858,6 +858,13 @@ def test_bam_dirmt(self, p):
 
         format.validate()
 
+        file_dict = format.file_dict(relative=True)
+
+        assert file_dict == {
+            "sample1": "sample1.bam",
+            "sample2": "sample2.bam",
+        }
+
     @patch('subprocess.run', return_value=Mock(returncode=3))
     def test_bam_dirmt_invalid(self, p):
         # this patch is not ideal but samtools' installation sometimes can
@@ -877,6 +884,19 @@ def test_multibam_dirmt(self, p):
 
         format.validate()
 
+        file_dict = format.file_dict(relative=True)
+
+        assert file_dict == {
+            "sample1": {
+                "map1": "sample1/map1.bam",
+                "map2": "sample1/map2.bam",
+            },
+            "sample2": {
+                "map1": "sample2/map1.bam",
+                "map2": "sample2/map2.bam",
+            },
+        }
+
 
 if __name__ == "__main__":
     unittest.main()