Coverage for src/mafw_tools/file_tools.py: 100%

1# Copyright 2025 European Union 

2# Author: Bulgheroni Antonio (antonio.bulgheroni@ec.europa.eu) 

3# SPDX-License-Identifier: EUPL-1.2 

4""" 

5File tools module. 

6 

7This module provides utilities for handling file operations, particularly 

8pickle serialization/deserialization and file name generation. 

9 

10.. versionadded:: 1.0.0 

11 

12Functions: 

13 save_to_pickle: Save dictionary-like objects to a pickle file 

14 load_from_pickle: Load objects from a pickle file 

15 generate_new_file_name: Generate a new file name based on an original file name 

16 

17Example usage:: 

18 

19 >>> from mafw_tools.file_tools import save_to_pickle, load_from_pickle 

20 >>> data = {'key': 'value'} 

21 >>> save_to_pickle('data.pkl', zipped=True, my_data=data) 

22 >>> loaded_data = load_from_pickle('data.pkl') 

23""" 

24import pickle 

25import zipfile 

26from pathlib import Path 

27from typing import Any 

28 

29 

30def save_to_pickle(filepath: str | Path, zipped: bool = True, **kwargs): 

31 """ 

32 Save a dictionary-like object to a pickle file. 

33 

34 This function is very helpful when there is the need to save multiple objects 

35 to the same binary file. The object to be serialized and saved to the file 

36 must be provided in the form of keyword arguments. 

37 

38 Here is an example: 

39 

40 .. code-block:: python 

41 

42 my_first_object = [1, 2, 3, 4] 

43 my_second_object = dict(field1='abc', field2='cba') 

44 

45 save_to_pickle( 

46 'my_pickle_file.sav', 

47 zipped=False, 

48 object1=my_first_object, 

49 object2=my_second_object, 

50 ) 

51 

52 A dictionary object in the form of :python:`{object1 : my_first_object, object2 : my_second_object}` 

53 is saved in the pickle file. 

54 

55 :param filepath: The output file name 

56 :type filepath: str | Path 

57 :param zipped: Flag to select if the output file should be compressed or not. Defaults to True 

58 :type zipped: bool, optional 

59 :param kwargs: The objects to the saved. 

60 """ 

61 objects_dict = kwargs 

62 

63 if isinstance(filepath, str): 

64 filepath = Path(filepath) 

65 

66 if zipped: 

67 # be sure about the extension 

68 if filepath.suffix == '.zip': 

69 filepath = filepath.with_suffix('.sav') 

70 

71 with open(filepath, 'wb') as f: 

72 pickle.dump(objects_dict, f) 

73 

74 if zipped: 

75 with zipfile.ZipFile(filepath.with_suffix('.zip'), 'w', compression=zipfile.ZIP_LZMA) as zip: 

76 zip.write(filepath, arcname=filepath.name) 

77 filepath.unlink() 

78 

79 

80def load_from_pickle(filepath: str | Path) -> dict[str, Any]: 

81 """ 

82 Load objects from a pickle file. 

83 

84 This function loads objects that were previously saved using :func:`save_to_pickle`. 

85 It handles both compressed and uncompressed pickle files automatically. 

86 

87 :param filepath: The path to the pickle file to load 

88 :type filepath: str | Path 

89 :return: Dictionary containing the loaded objects 

90 :rtype: dict[str, Any] 

91 """ 

92 if isinstance(filepath, str): 

93 filepath = Path(filepath) 

94 

95 zipped = zipfile.is_zipfile(filepath) 

96 

97 if zipped: 

98 with zipfile.ZipFile(filepath) as zip: 

99 with zip.open(filepath.with_suffix('.sav').name) as file: 

100 objects_dict = pickle.load(file) 

101 else: 

102 with open(filepath, 'rb') as file: 

103 objects_dict = pickle.load(file) 

104 

105 return objects_dict 

106 

107 

108def generate_new_file_name( 

109 original_file_name: str | Path, new_base_path: str | Path, extra_suffix: str = None, new_extension: str = None 

110) -> Path: 

111 """ 

112 Generate a new file name including its full path based on an original file name. 

113 

114 The generated file name will have the new base path provided, an extra suffix 

115 (if provided) just before the extension, and if a new extension is provided 

116 then this is also replaced. 

117 

118 :param original_file_name: The starting file name. It can be the full file name of an image. It does not need to be the full path. The filename is just fine. 

119 :type original_file_name: str | Path 

120 :param new_base_path: The new base path to be appended to the newly generated filename. 

121 :type new_base_path: str | Path 

122 :param extra_suffix: A string to be added just before the file extension. 

123 :type extra_suffix: str, optional, defaults to None 

124 :param new_extension: A new extension to replace the old one. 

125 :type new_extension: str, optional, defaults to None 

126 :return: The newly generated file name 

127 :rtype: Path 

128 """ 

129 

130 if isinstance(original_file_name, str): 

131 original_file_name = Path(original_file_name) 

132 

133 if isinstance(new_base_path, str): 

134 new_base_path = Path(new_base_path) 

135 

136 # be sure that the new path exists. 

137 new_base_path.mkdir(exist_ok=True, parents=True) 

138 

139 if new_extension: 

140 if not new_extension.startswith('.'): 

141 new_extension = '.' + new_extension 

142 new_file_name = Path(original_file_name.with_suffix(new_extension).name) 

143 else: 

144 new_file_name = Path(original_file_name.name) 

145 

146 if extra_suffix: 

147 if not extra_suffix.startswith('_'): 

148 extra_suffix = '_' + extra_suffix 

149 new_file_name = Path(new_file_name.stem + extra_suffix + new_file_name.suffix) 

150 

151 return new_base_path / new_file_name