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

1 import shutil

2 from pathlib import Path

3 import random

4 import string

5 from datetime import datetime

6 import pickle

7

8

9

10 def delete_create_folder(path: Path):

11 shutil.rmtree(path, ignore_errors=True)

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

13

14

15 def get_rand_hex():

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

17

18

19 def get_date_time_str():

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

21

22

23 def strip_unsafe(filename):

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

25

26

27 def save_pickle(file_path: Path):

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

29 return pickle.load(f)

30

31

32 def load_pickle(obj, file_path: Path):

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

34 pickle.dump(obj, f)

35

36

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

38

39

40 def 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

54 def 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)