updated get_format function

locidex/constants.py

@@ -64,8 +64,8 @@ def _keys(cls) -> list:
         return [i.name for i in fields(cls)]
 
 FILE_TYPES = {
-    'genbank': [".gbk",".genbank",".gbf",".gbk.gz",".genbank.gz",".gbf.gz",".gbff",".gbff.gz"],
-    'fasta': [".fasta",".fas",".fa",".ffn",".fna",".fasta.gz",".fas.gz",".fa.gz",".ffn.gz",".fna.gz"],
+    'genbank': frozenset([".gbk",".genbank",".gbf",".gbk.gz",".genbank.gz",".gbf.gz",".gbff",".gbff.gz"]),
+    'fasta': frozenset([".fasta",".fas",".fa",".ffn",".fna",".fasta.gz",".fas.gz",".fa.gz",".ffn.gz",".fna.gz"]),
 }
 
 

locidex/extract.py

@@ -18,7 +18,7 @@
 from locidex.constants import SEARCH_RUN_DATA, FILE_TYPES, BlastColumns, BlastCommands, DBConfig, DB_EXPECTED_FILES, EXTRACT_MODES, raise_file_not_found_e
 from locidex.version import __version__
 from locidex.classes.aligner import perform_alignment, aligner
-from locidex.utils import check_db_groups
+from locidex.utils import check_db_groups, get_format
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(filemode=sys.stderr, level=logging.INFO)
@@ -112,10 +112,7 @@ def run_extract(config):
 
     format = None
     if format is None:
-        for t in FILE_TYPES:
-            for ext in FILE_TYPES[t]:
-                if ext == input_fasta.suffix:
-                    format = t
+        format = get_format(input_fasta)
     else:
         format = format.lower()
 

locidex/format.py

@@ -238,7 +238,7 @@ def run(cmd_args=None):
 
     logger.info("Beginning format operation.")
     obj = locidex_format(input=input,header=LocidexDBHeader._fields,is_protein=is_coding,min_len_frac=min_len_frac,max_len_frac=max_len_frac, min_ident_perc=min_ident,
-            min_cov_perc=min_match_cov,trans_table=trans_table,valid_ext=FILE_TYPES['fasta'])
+            min_cov_perc=min_match_cov,trans_table=trans_table,valid_ext=list(FILE_TYPES['fasta']))
     logger.info("Finished format.")
     run_data['result_file'] = os.path.join(outdir,"locidex.txt")
     pd.DataFrame.from_dict(obj.data,orient='index').to_csv(run_data['result_file'],sep="\t",index=False,header=True)

locidex/search.py

@@ -16,7 +16,7 @@
 from locidex.manifest import DBData
 from locidex.classes.seq_intake import seq_intake, seq_store, HitFilters
 from locidex.constants import BlastCommands,  SEARCH_RUN_DATA, FILE_TYPES, BlastColumns, DB_EXPECTED_FILES, OPTION_GROUPS, DBConfig
-from locidex.utils import write_seq_dict, check_db_groups, slots
+from locidex.utils import write_seq_dict, check_db_groups, slots, get_format
 from locidex.version import __version__
 
 logger = logging.getLogger(__name__)
@@ -151,10 +151,7 @@ def run_search(config):
         os.makedirs(outdir, 0o755)
 
     if format is None:
-        for t in FILE_TYPES:
-            for ext in FILE_TYPES[t]:
-                if query_file.suffix == ext:
-                    format = t
+        format = get_format(query_file)
     else:
         format = format.lower()
 

locidex/utils.py

@@ -6,8 +6,8 @@
 from pathlib import Path
 from locidex.manifest import ManifestItem
 from Bio.Seq import Seq
-from typing import Dict, FrozenSet
-from locidex.constants import NT_SUB, PROTEIN_ALPHA, DNA_ALPHA, OPTION_GROUPS
+from typing import Dict, FrozenSet, Optional
+from locidex.constants import NT_SUB, PROTEIN_ALPHA, DNA_ALPHA, OPTION_GROUPS, FILE_TYPES
 import locidex.manifest as manifest 
 
 def slots(annotations: Dict[str, object]) -> FrozenSet[str]:
@@ -150,3 +150,15 @@ def filter_hsps_df(df):
     return
 
 
+def get_format(file: Path) -> Optional[str]:
+    """
+    Return file type based on suffix used
+    """
+    format: str = None
+    file_exts = file.suffixes
+    for k, extensions in FILE_TYPES.items():
+        for ext in file_exts:
+            if ext in extensions:
+                format = k
+                break
+    return format

tests/test_utils.py

@@ -4,6 +4,7 @@
 
 import pytest
 from locidex import utils
+from pathlib import Path
 from locidex import manifest
 from argparse import Namespace
 from collections import namedtuple
@@ -26,4 +27,32 @@ def test_check_db_groups_fail():
     analysis_params = {"db_group": "Db1", "db_name": "test_name"}
 
     with pytest.raises(KeyError):
-        analysis_params = utils.check_db_groups(analysis_params, nm_group)
+        analysis_params = utils.check_db_groups(analysis_params, nm_group)
+
+
+@pytest.mark.parametrize( "file_in,type",
+    [
+        ("test.fa", "fasta"),
+        ("test.fas", "fasta"),
+        ("test.ffn", "fasta"),
+        ("test.fna", "fasta"),
+        ("test.fasta.gz", "fasta"),
+        ("test.fas.gz", "fasta"),
+        ("test.fa.gz", "fasta"),
+        ("test.fna.gz", "fasta"),
+        ("test.ffn.gz", "fasta"),
+        ("test.gbk", "genbank"),
+        ("test.genbank", "genbank"),
+        ("test.gbf", "genbank"),
+        ("test.gbk.gz", "genbank"),
+        ("test.genbank.gz", "genbank"),
+        ("test.gbf.gz", "genbank"),
+        ("test.gbff.gz", "genbank"),
+        ("test.gbff", "genbank"),
+        ("test.1.fasta.gz", "fasta"),
+    ])
+def test_get_format(file_in, type):
+    """
+    test get_format function
+    """
+    assert utils.get_format(Path(file_in)) == type