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"""

5 File tools module.

6

7 This module provides utilities for handling file operations, particularly

8 pickle serialization/deserialization and file name generation.

9

10.. versionadded:: 1.0.0

11

12 Functions:

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

17 Example 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"""

24 import pickle

25 import zipfile

26 from pathlib import Path

27 from typing import Any

28

29

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

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

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