added file_dict mixin

Author: VinzentRisch

q2_types/kraken2/_formats.py

@@ -5,6 +5,7 @@
 #
 # The full license is in the file LICENSE, distributed with this software.
 # ----------------------------------------------------------------------------
+from collections import defaultdict
 
 import pandas as pd
 from pandas.core.dtypes.common import is_string_dtype
@@ -67,7 +68,104 @@ def _validate_(self, level):
             )
 
 
-class Kraken2ReportDirectoryFormat(model.DirectoryFormat):
+class FileDictMixin:
+    def file_dict(self, relative=False, suffixes=None):
+        """
+        For per sample directories it returns a mapping of sample id to
+        another dictionary where keys represent the file name and values
+        correspond to the filepath for each file.
+        For files, it returns a mapping of file name to filepath for each
+        file. The specified suffixes are removed from filenames.
+
+        Parameters
+        ---------
+        relative : bool
+            Whether to return filepaths relative to the directory's location.
+            Returns absolute filepaths by default.
+        suffixes : List
+            A list of suffixes that should be removed from the filenames to
+            generate the ID.
+
+        Returns
+        -------
+        dict
+            Mapping of filename -> filepath as described above.
+            Or mapping of sample id -> dict {filename: filepath} as
+            described above.
+            Both levels of the dictionary are sorted alphabetically by key.
+        """
+        ids = defaultdict(dict)
+        for entry in self.path.iterdir():
+            if entry.is_dir():
+                outer_id = entry.name
+                for path in entry.iterdir():
+                    file_path, inner_id = _create_path(
+                        path=path,
+                        relative=relative,
+                        dir_format=self,
+                        suffixes=suffixes
+                    )
+
+                    ids[outer_id][inner_id] = str(file_path)
+                ids[outer_id] = dict(sorted(ids[outer_id].items()))
+            else:
+                file_path, inner_id = _create_path(
+                    path=entry,
+                    relative=relative,
+                    dir_format=self,
+                    suffixes=suffixes
+
+                )
+
+                ids[inner_id] = str(file_path)
+
+        return dict(sorted(ids.items()))
+
+
+def _create_path(path, relative, dir_format, suffixes):
+    """
+    This function processes the input file path to generate an absolute or
+    relative path string and the ID derived from the file name. The ID is
+    extracted by removing the one of the specified suffixes  from the file
+    name. If no suffixes are specified the ID is defined to be the filename.
+
+    Parameters:
+    ---------
+        path : Path
+            A Path object representing the file path to process.
+        relative : bool
+            A flag indicating whether the returned path should be relative
+            to the directory formats path or absolute.
+        dir_format : DirectoryFormat.
+            Any object of class DirectoryFormat.
+
+    Returns:
+    -------
+        path_dict : str
+            The full relative or absolut path to the file.
+        _id : str
+            The ID derived from the file name. ID will be "" if the filename
+            consists only of the suffix.
+    """
+    file_name = path.stem
+
+    _id = file_name
+
+    if suffixes:
+        for suffix in suffixes:
+            if file_name.endswith(suffix[1:]):
+                _id = file_name[:-len(suffix)]
+                break
+
+    path_dict = (
+        path.absolute().relative_to(dir_format.path.absolute())
+        if relative
+        else path.absolute()
+    )
+    return str(path_dict), _id
+
+
+class Kraken2ReportDirectoryFormat(model.DirectoryFormat, FileDictMixin):
     reports = model.FileCollection(
         r'.+report\.(txt|tsv)$', format=Kraken2ReportFormat
     )

q2_types/kraken2/tests/test_formats.py

@@ -87,6 +87,12 @@ def test_report_dirfmt_from_reads(self):
         fmt = Kraken2ReportDirectoryFormat(dirpath, mode='r')
         fmt.validate()
 
+    def test_report_dirfmt_from_reads_dict(self):
+        dirpath = self.get_data_path('reports-reads')
+        fmt = Kraken2ReportDirectoryFormat(dirpath, mode='r')
+        dict = fmt.file_dict(suffixes=["_report"])
+        print(fmt)
+
     def test_report_dirfmt_from_mags(self):
         dirpath = self.get_data_path('reports-mags')
         fmt = Kraken2ReportDirectoryFormat(dirpath, mode='r')