Coverage for b4_backup/utils.py: 100%

1"""A collection of Helper functions.""" 

2 

3import logging 

4import os 

5from pathlib import Path, PurePath 

6 

7from omegaconf import OmegaConf, SCMode 

8from rich.console import Console 

9from rich.logging import RichHandler 

10from rich.theme import Theme 

11 

12from b4_backup.config_schema import DEFAULT, BaseConfig 

13 

14log = logging.getLogger("b4_backup.utils") 

15 

16DEFAULT_CONFIG = Path(os.getenv("B4_BACKUP_CONFIG", str(BaseConfig.config_path))) 

17 

18DEFAULT_THEME = Theme( 

19 { 

20 "logging.level.debug": "blue", 

21 "logging.level.info": "green", 

22 "logging.level.warning": "yellow", 

23 "logging.level.error": "red", 

24 "logging.level.critical": "reverse red", 

25 } 

26) 

27CONSOLE = Console(theme=DEFAULT_THEME) 

28 

29 

30# Dynamically imported by logging.config.dictConfig 

31def rich_handler() -> RichHandler: 

32 """Used in the logging config to use a customized RichHandler.""" 

33 return RichHandler(console=CONSOLE) 

34 

35 

36def resolve_parent_dir(path: str) -> str: 

37 """ 

38 This resolver (parent_dir) can be used in OmegaConf configs to return the parent directory or the directory of a file. 

39 

40 Args: 

41 path: Name of the file or directory 

42 

43 Returns: 

44 Parent directory 

45 """ 

46 return str(PurePath(path).parent) 

47 

48 

49def resolve_from_file(path: str) -> str: 

50 """ 

51 This resolver (from_file) can be used in OmegaConf configs to use the raw content of a file as an input. 

52 

53 Args: 

54 path: Input file path 

55 

56 Returns: 

57 File content 

58 """ 

59 return Path(path).read_text(encoding="utf8").strip() 

60 

61 

62def _copy_from_default_retention(config: BaseConfig): 

63 for target in config.backup_targets.values(): 

64 for ( 

65 retention_name, 

66 retention, 

67 ) in config.backup_targets[DEFAULT].src_retention.items(): 

68 if retention_name not in target.src_retention: 

69 target.src_retention[retention_name] = retention 

70 

71 for ( 

72 retention_name, 

73 retention, 

74 ) in config.backup_targets[DEFAULT].dst_retention.items(): 

75 if retention_name not in target.dst_retention: 

76 target.dst_retention[retention_name] = retention 

77 

78 

79def load_config( 

80 config_path: Path = DEFAULT_CONFIG, overrides: list[str] | None = None 

81) -> BaseConfig: 

82 """ 

83 Reads the config file and returns a config dataclass. 

84 

85 Args: 

86 config_path: Path of the config file 

87 overrides: 

88 A list of dot list entries, which can override the values in the config. 

89 Used for CLI. 

90 

91 Returns: 

92 Config object 

93 """ 

94 overrides = overrides or [] 

95 

96 config_path = config_path.expanduser() 

97 config_path.parent.mkdir(exist_ok=True, parents=True) 

98 _ = config_path.exists() or config_path.touch() 

99 

100 if not OmegaConf.has_resolver("from_file"): 

101 OmegaConf.register_new_resolver("from_file", resolve_from_file) 

102 OmegaConf.register_new_resolver("parent_dir", resolve_parent_dir) 

103 

104 base_conf = OmegaConf.merge( 

105 OmegaConf.structured(BaseConfig), 

106 OmegaConf.load(config_path), 

107 OmegaConf.from_dotlist(overrides), 

108 ) 

109 

110 # Templates shouldn't fail, if there is a value missing 

111 base_conf.backup_targets[DEFAULT].source = "NONE" 

112 

113 # pylance doesn't understand here that it's actually a config_schema.BaseConfig type 

114 base_conf_instance: BaseConfig = OmegaConf.to_container( # type: ignore 

115 base_conf, structured_config_mode=SCMode.INSTANTIATE, resolve=True 

116 ) 

117 

118 base_conf_instance.config_path = config_path 

119 

120 # retention rulesets shouldn't do a nested merge 

121 # That's why I do a shallow update here manually 

122 _copy_from_default_retention(base_conf_instance) 

123 

124 return base_conf_instance 

125 

126 

127def contains_path(path: PurePath, sub_path: PurePath) -> bool: 

128 """ 

129 Check if a subpath is included in another path. 

130 

131 Args: 

132 path: Path you want to check 

133 sub_path: Subpath that should be included in path. 

134 

135 Returns: 

136 True if path contains subpath 

137 """ 

138 return any( 

139 slice == sub_path.parts 

140 for slice in zip(*[path.parts[i:] for i in range(len(sub_path.parts))]) 

141 )