local_files_collector.py 2.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. import os
  2. from tqdm.auto import tqdm
  3. class Local_Files(object):
  4. """Local_Files gives all the wanted files in dictionary from the given path
  5. Maintaining file_type_dicts states which format belongs to which category.
  6. file_type_dicts Can be uploaded using class methods update_file_type_dicts,add_new_itm_file_type_dicts
  7. """
  8. file_type_dicts = {
  9. "images": ["jpg", "png", "svg", "webp"],
  10. "media_files": ["mp3", "mp4", "mpeg4", "wmv", "3gp", "webm"],
  11. "documents": ["doc", "docx", "csv", "pdf", ]
  12. }
  13. def __init__(self, root_dir=None, new_file_type_dicts=None):
  14. """Root path is mandatory to read files from root path and its sub directory
  15. """
  16. if root_dir:
  17. if os.path.exists(root_dir):
  18. self.root_dir = root_dir
  19. if new_file_type_dicts:
  20. self.result = self.get_files_dict(new_file_type_dicts)
  21. print("Root Directory set to :\t{}".format(self.root_dir))
  22. self.result = self.get_files_dict()
  23. else:
  24. raise Exception('Root Directory not found')
  25. def get_files_dict(self, file_type_dict=None, file_type_dicts=None):
  26. """get_files_dict iterate over all the folders and its sub directory,
  27. get_files_dict collects all the desired files and create dictionary which is useful in uploading files"""
  28. if file_type_dict is None:
  29. file_type_dict = file_type_dicts
  30. file_wit_key_value = {}
  31. for root, dirs, files in tqdm(os.walk(self.root_dir)):
  32. for each_file in files:
  33. for each_key, each_item in file_type_dict.items():
  34. if each_file.rsplit(".")[-1] in each_item:
  35. if each_key in file_wit_key_value.keys():
  36. file_wit_key_value[each_key].append(
  37. os.path.join(root, each_file))
  38. else:
  39. file_wit_key_value[each_key] = [
  40. os.path.join(root, each_file)]
  41. return file_wit_key_value
  42. @classmethod
  43. def udpate_file_type_dicts(cls, key, value):
  44. """Class method to update values in default dictionary
  45. Example for document you want to add txt file this method can be useful"""
  46. if key in cls.file_type_dicts.items():
  47. cls.file_type_dicts[key].append(value)
  48. @classmethod
  49. def add_new_itm_file_type_dicts(cls, key, value):
  50. if key not in cls.file_type_dicts.items():
  51. cls.file_type_dicts[key] = value