Coverage for src/image_utils/file_utils.py: 80%

35 statements  

« prev     ^ index     » next       coverage.py v7.2.5, created at 2023-05-19 16:15 -0700

1import shutil 

2from pathlib import Path 

3import random 

4import string 

5from datetime import datetime 

6import pickle 

7 

8 

9 

10def delete_create_folder(path: Path): 

11 shutil.rmtree(path, ignore_errors=True) 

12 path.mkdir(parents=True, exist_ok=True) 

13 

14 

15def get_rand_hex(): 

16 return ''.join(random.choices(string.ascii_uppercase + string.digits, k=5)) 

17 

18 

19def get_date_time_str(): 

20 return datetime.now().strftime("%Y_%m_%d-%H_%M") 

21 

22 

23def strip_unsafe(filename): 

24 return "".join([c for c in filename if c.isalpha() or c.isdigit() or c == '' or c == '_']).rstrip() 

25 

26 

27def save_pickle(file_path: Path): 

28 with open(file_path, "rb") as f: 

29 return pickle.load(f) 

30 

31 

32def load_pickle(obj, file_path: Path): 

33 with open(file_path, "wb") as f: 

34 pickle.dump(obj, f) 

35 

36 

37def get_file_list(**kwargs): [f for f in get_files(**kwargs)] 

38 

39 

40def get_files(path: Path, recursive: bool = False, return_folders: bool = False, allowed_extensions=None): 

41 path = Path(path) 

42 

43 if allowed_extensions or recursive: 

44 glob_str = "*" if allowed_extensions is None else f"*.[{''.join(allowed_extensions)}]*" 

45 iterator = path.rglob(glob_str) if recursive else path.glob(glob_str) 

46 else: 

47 iterator = path.iterdir() 

48 

49 for file in iterator: 

50 if file.is_file() or return_folders: 

51 yield file 

52 

53 

54def get_images(recursive: bool = False, allowed_extensions=[".png", ".jpg", ".jpeg"]): 

55 from image_utils import Im 

56 for file in get_files(recursive=recursive, allowed_extensions=allowed_extensions): 

57 yield Im.open(file)