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
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-20 10:58 +0200
1"""Mixin for cache with joblib."""
3import os
4import warnings
5from pathlib import Path
7from joblib import Memory
9import nilearn
10from nilearn._utils.logger import find_stack_level
12from .helpers import stringify_path
14MEMORY_CLASSES = (Memory,)
17def check_memory(memory, verbose=0):
18 """Ensure an instance of a joblib.Memory object.
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.
26 verbose : int, default=0
27 Verbosity level.
29 Returns
30 -------
31 memory : instance of joblib.Memory.
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()
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)
74 memory = Memory(location=str(cache_dir), verbose=verbose)
75 return memory
78class _ShelvedFunc:
79 """Work around for Python 2, for which pickle fails on instance method."""
81 def __init__(self, func):
82 self.func = func
83 self.func_name = f"{func.__name__}_shelved"
85 def __call__(self, *args, **kwargs):
86 return self.func.call_and_shelve(*args, **kwargs)
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.
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.
105 Parameters
106 ----------
107 func : function
108 The function which output is to be cached.
110 memory : instance of joblib.Memory, string or pathlib.Path
111 Used to cache the function call.
113 func_memory_level : int, optional
114 The memory_level from which caching must be enabled for the wrapped
115 function.
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).
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.
126 kwargs : keyword arguments, optional
127 The keyword arguments passed to memory.cache.
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.
137 """
138 verbose = kwargs.get("verbose", 0)
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)
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 )
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
183class CacheMixin:
184 """Mixin to add caching to a class.
186 This class is a thin layer on top of joblib.Memory, that mainly adds a
187 "caching level", similar to a "log level".
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.
196 """
198 def _cache(self, func, func_memory_level=1, shelve=False, **kwargs):
199 """Return a joblib.Memory object.
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.
207 Parameters
208 ----------
209 func : function
210 The function the output of which is to be cached.
212 func_memory_level : int, default=1
213 The memory_level from which caching must be enabled for the wrapped
214 function.
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.
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.
228 """
229 verbose = getattr(self, "verbose", 0)
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)
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
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 )