Coverage for src/appl/core/io.py: 67%

1import json 

2import os 

3from typing import IO, Any, Callable, Dict, Optional 

4 

5import addict 

6import toml 

7import yaml 

8 

9# from importlib import import_module # read python 

10PLAIN_TEXT_FILES = [".txt", ".log", ".md", ".html"] 

11 

12 

13def makedirs(file: str) -> None: 

14 """Make the directory of the file if it does not exist.""" 

15 if folder := os.path.dirname(file): 

16 os.makedirs(folder, exist_ok=True) 

17 

18 

19def get_ext(file: str) -> str: 

20 """Get the extension of a file.""" 

21 return os.path.splitext(file)[1] 

22 

23 

24def dump_file( 

25 data: Any, 

26 file: str, 

27 mode: str = "w", 

28 ensure_folder_exists: bool = True, 

29 file_type: Optional[str] = None, 

30 *args: Any, 

31 **kwargs: Any, 

32) -> None: 

33 """Write the data to a file based on the file extension.""" 

34 if file_type is None: 

35 file_type = get_ext(file) 

36 if ensure_folder_exists: 

37 makedirs(file) 

38 

39 if file_type == ".json": 

40 dump_func: Callable = json.dump 

41 elif file_type in [".yaml", ".yml"]: 

42 dump_func = yaml.dump 

43 elif file_type == ".toml": 

44 dump_func = toml.dump 

45 elif file_type in PLAIN_TEXT_FILES: 

46 

47 def dump_func(data, f, *args, **kwargs): 

48 f.write(data) 

49 

50 else: 

51 raise ValueError(f"Unsupported file type {file_type}") 

52 with open(file, mode) as f: 

53 dump_func(data, f, *args, **kwargs) 

54 

55 

56def load_file( 

57 file: str, 

58 mode: str = "r", 

59 file_type: Optional[str] = None, 

60 open_kwargs: Optional[Dict[str, Any]] = None, 

61 *args: Any, 

62 **kwargs: Any, 

63) -> Any: 

64 """Load a file based on the file extension and return the data.""" 

65 if file_type is None: 

66 file_type = get_ext(file) 

67 if file_type == ".json": 

68 load_func: Callable = json.load 

69 elif file_type in [".yaml", ".yml"]: 

70 load_func = yaml.safe_load 

71 elif file_type == ".toml": 

72 load_func = toml.load 

73 # elif file_type == ".py": 

74 # load_func = import_module 

75 elif file_type in PLAIN_TEXT_FILES: 

76 

77 def load_func(f: IO[Any], *args: Any, **kwargs: Any) -> Any: 

78 return f.read() 

79 

80 else: 

81 raise ValueError(f"Unsupported file type {file_type}") 

82 open_kwargs = open_kwargs or {} 

83 with open(file, mode, **open_kwargs) as f: 

84 return load_func(f, *args, **kwargs)