Coverage for nilearn/_utils/cache_mixin.py: 16%

64 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-20 10:58 +0200

1"""Mixin for cache with joblib.""" 

2 

3import os 

4import warnings 

5from pathlib import Path 

6 

7from joblib import Memory 

8 

9import nilearn 

10from nilearn._utils.logger import find_stack_level 

11 

12from .helpers import stringify_path 

13 

14MEMORY_CLASSES = (Memory,) 

15 

16 

17def check_memory(memory, verbose=0): 

18 """Ensure an instance of a joblib.Memory object. 

19 

20 Parameters 

21 ---------- 

22 memory : None, instance of joblib.Memory, str or pathlib.Path 

23 Used to cache the masking process. 

24 If a str is given, it is the path to the caching directory. 

25 

26 verbose : int, default=0 

27 Verbosity level. 

28 

29 Returns 

30 ------- 

31 memory : instance of joblib.Memory. 

32 

33 """ 

34 if memory is None: 

35 memory = Memory(location=None, verbose=verbose) 

36 # TODO make Path the default here 

37 memory = stringify_path(memory) 

38 if isinstance(memory, str): 

39 cache_dir = memory 

40 if nilearn.EXPAND_PATH_WILDCARDS: 

41 cache_dir = Path(cache_dir).expanduser() 

42 

43 # Perform some verifications on given path. 

44 split_cache_dir = os.path.split(cache_dir) 

45 if len(split_cache_dir) > 1 and ( 

46 not Path(split_cache_dir[0]).exists() and split_cache_dir[0] != "" 

47 ): 

48 if not nilearn.EXPAND_PATH_WILDCARDS and cache_dir.startswith("~"): 

49 # Maybe the user want to enable expanded user path. 

50 error_msg = ( 

51 "Given cache path parent directory doesn't " 

52 f"exists, you gave '{split_cache_dir[0]}'. Enabling " 

53 "nilearn.EXPAND_PATH_WILDCARDS could solve " 

54 "this issue." 

55 ) 

56 elif memory.startswith("~"): 

57 # Path built on top of expanded user path doesn't exist. 

58 error_msg = ( 

59 "Given cache path parent directory doesn't " 

60 f"exists, you gave '{split_cache_dir[0]}' " 

61 "which was expanded as '{os.path.dirname(memory)}' " 

62 "but doesn't exist either. " 

63 "Use nilearn.EXPAND_PATH_WILDCARDS to deactivate " 

64 "auto expand user path (~) behavior." 

65 ) 

66 else: 

67 # The given cache base path doesn't exist. 

68 error_msg = ( 

69 "Given cache path parent directory doesn't " 

70 "exists, you gave '{split_cache_dir[0]}'." 

71 ) 

72 raise ValueError(error_msg) 

73 

74 memory = Memory(location=str(cache_dir), verbose=verbose) 

75 return memory 

76 

77 

78class _ShelvedFunc: 

79 """Work around for Python 2, for which pickle fails on instance method.""" 

80 

81 def __init__(self, func): 

82 self.func = func 

83 self.func_name = f"{func.__name__}_shelved" 

84 

85 def __call__(self, *args, **kwargs): 

86 return self.func.call_and_shelve(*args, **kwargs) 

87 

88 

89def cache( 

90 func, 

91 memory, 

92 func_memory_level=None, 

93 memory_level=None, 

94 shelve=False, 

95 **kwargs, 

96): 

97 """Return a joblib.Memory object. 

98 

99 The memory_level determines the level above which the wrapped 

100 function output is cached. By specifying a numeric value for 

101 this level, the user can to control the amount of cache memory 

102 used. This function will cache the function call or not 

103 depending on the cache level. 

104 

105 Parameters 

106 ---------- 

107 func : function 

108 The function which output is to be cached. 

109 

110 memory : instance of joblib.Memory, string or pathlib.Path 

111 Used to cache the function call. 

112 

113 func_memory_level : int, optional 

114 The memory_level from which caching must be enabled for the wrapped 

115 function. 

116 

117 memory_level : int, optional 

118 The memory_level used to determine if function call must 

119 be cached or not (if user_memory_level is equal of greater than 

120 func_memory_level the function is cached). 

121 

122 shelve : bool, default=False 

123 Whether to return a joblib MemorizedResult, callable by a .get() 

124 method, instead of the return value of func. 

125 

126 kwargs : keyword arguments, optional 

127 The keyword arguments passed to memory.cache. 

128 

129 Returns 

130 ------- 

131 mem : joblib.MemorizedFunc, wrapped in _ShelvedFunc if shelving 

132 Object that wraps the function func to cache its further call. 

133 This object may be a no-op, if the requested level is lower 

134 than the value given to _cache()). 

135 For consistency, a callable object is always returned. 

136 

137 """ 

138 verbose = kwargs.get("verbose", 0) 

139 

140 # memory_level and func_memory_level must be both None or both integers. 

141 memory_levels = [memory_level, func_memory_level] 

142 both_params_integers = all(isinstance(lvl, int) for lvl in memory_levels) 

143 both_params_none = all(lvl is None for lvl in memory_levels) 

144 

145 if not (both_params_integers or both_params_none): 

146 raise ValueError( 

147 "Reference and user memory levels must be both None " 

148 "or both integers." 

149 ) 

150 

151 if memory is not None and ( 

152 func_memory_level is None or memory_level >= func_memory_level 

153 ): 

154 memory = stringify_path(memory) 

155 if isinstance(memory, str): 

156 memory = Memory(location=memory, verbose=verbose) 

157 if not isinstance(memory, MEMORY_CLASSES): 

158 raise TypeError( 

159 "'memory' argument must be a string or a " 

160 "joblib.Memory object. " 

161 f"{memory} {type(memory)} was given." 

162 ) 

163 if ( 

164 memory.location is None 

165 and memory_level is not None 

166 and memory_level > 1 

167 ): 

168 warnings.warn( 

169 f"Caching has been enabled (memory_level = {memory_level}) " 

170 "but no Memory object or path has been provided" 

171 " (parameter memory). Caching deactivated for " 

172 f"function {func.__name__}.", 

173 stacklevel=find_stack_level(), 

174 ) 

175 else: 

176 memory = Memory(location=None, verbose=verbose) 

177 cached_func = memory.cache(func, **kwargs) 

178 if shelve: 

179 cached_func = _ShelvedFunc(cached_func) 

180 return cached_func 

181 

182 

183class CacheMixin: 

184 """Mixin to add caching to a class. 

185 

186 This class is a thin layer on top of joblib.Memory, that mainly adds a 

187 "caching level", similar to a "log level". 

188 

189 Notes 

190 ----- 

191 Usage: to cache the results of a method, wrap it in self._cache() 

192 defined by this class. Caching is performed only if the user-specified 

193 cache level (self._memory_level) is greater than the value given as a 

194 parameter to self._cache(). See _cache() documentation for details. 

195 

196 """ 

197 

198 def _cache(self, func, func_memory_level=1, shelve=False, **kwargs): 

199 """Return a joblib.Memory object. 

200 

201 The memory_level determines the level above which the wrapped 

202 function output is cached. By specifying a numeric value for 

203 this level, the user can to control the amount of cache memory 

204 used. This function will cache the function call or not 

205 depending on the cache level. 

206 

207 Parameters 

208 ---------- 

209 func : function 

210 The function the output of which is to be cached. 

211 

212 func_memory_level : int, default=1 

213 The memory_level from which caching must be enabled for the wrapped 

214 function. 

215 

216 shelve : bool, default=False 

217 Whether to return a joblib MemorizedResult, callable by a .get() 

218 method, instead of the return value of func. 

219 

220 Returns 

221 ------- 

222 mem : joblib.MemorizedFunc, wrapped in _ShelvedFunc if shelving 

223 Object that wraps the function func to cache its further call. 

224 This object may be a no-op, if the requested level is lower 

225 than the value given to _cache()). 

226 For consistency, a callable object is always returned. 

227 

228 """ 

229 verbose = getattr(self, "verbose", 0) 

230 

231 # Creates attributes if they don't exist 

232 # This is to make creating them in __init__() optional. 

233 if not hasattr(self, "memory_level"): 

234 self.memory_level = 0 

235 if not hasattr(self, "memory"): 

236 self.memory = Memory(location=None, verbose=verbose) 

237 self.memory = check_memory(self.memory, verbose=verbose) 

238 

239 # If cache level is 0 but a memory object has been provided, set 

240 # memory_level to 1 with a warning. 

241 if self.memory_level == 0 and self.memory.location is not None: 

242 warnings.warn( 

243 "memory_level is currently set to 0 but " 

244 "a Memory object has been provided. " 

245 "Setting memory_level to 1.", 

246 stacklevel=find_stack_level(), 

247 ) 

248 self.memory_level = 1 

249 

250 return cache( 

251 func, 

252 self.memory, 

253 func_memory_level=func_memory_level, 

254 memory_level=self.memory_level, 

255 shelve=shelve, 

256 **kwargs, 

257 )