Переглянути джерело

refactor get_all_filenames

quarrying 4 роки тому
батько
коміт
5f10fe7d64
1 змінених файлів з 16 додано та 0 видалено
  1. 16 0
      khandy/utils_fs.py

+ 16 - 0
khandy/utils_fs.py

@@ -79,12 +79,28 @@ def listdirs(paths, path_sep=None, full_path=True):
     return all_filenames
 
 
+def _normalize_extensions(extensions):
+    if extensions is None:
+        return None
+    new_extensions = []
+    
+    for extension in extensions:
+        if extension.startswith('.'):
+            new_extensions.append(extension.lower())
+        else:
+            new_extensions.append('.' + extension.lower())
+    return new_extensions
+
+
 def get_all_filenames(path, extensions=None, is_valid_file=None):
     if (extensions is not None) and (is_valid_file is not None):
         raise ValueError("Both extensions and is_valid_file cannot "
                          "be not None at the same time")
     if is_valid_file is None:
         if extensions is not None:
+            if isinstance(extensions, str):
+                extensions = [extensions]
+            extensions = tuple(_normalize_extensions(extensions))
             def is_valid_file(filename):
                 return filename.lower().endswith(extensions)
         else: