Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""Utilities for fast persistence of big data, with optional compression."""
3# Author: Gael Varoquaux <gael dot varoquaux at normalesup dot org>
4# Copyright (c) 2009 Gael Varoquaux
5# License: BSD Style, 3 clauses.
7import contextlib
8import io
9import pickle
10import warnings
12from .compressor import _COMPRESSORS, _ZFILE_PREFIX
14try:
15 import numpy as np
16except ImportError:
17 np = None
19Unpickler = pickle._Unpickler
20Pickler = pickle._Pickler
21xrange = range
24try:
25 # The python standard library can be built without bz2 so we make bz2
26 # usage optional.
27 # see https://github.com/scikit-learn/scikit-learn/issues/7526 for more
28 # details.
29 import bz2
30except ImportError:
31 bz2 = None
33# Buffer size used in io.BufferedReader and io.BufferedWriter
34_IO_BUFFER_SIZE = 1024**2
37def _is_raw_file(fileobj):
38 """Check if fileobj is a raw file object, e.g created with open."""
39 fileobj = getattr(fileobj, "raw", fileobj)
40 return isinstance(fileobj, io.FileIO)
43def _get_prefixes_max_len():
44 # Compute the max prefix len of registered compressors.
45 prefixes = [len(compressor.prefix) for compressor in _COMPRESSORS.values()]
46 prefixes += [len(_ZFILE_PREFIX)]
47 return max(prefixes)
50def _ensure_native_byte_order(array):
51 """Use the byte order of the host while preserving values
53 Does nothing if array already uses the system byte order.
54 """
55 native_dtype = array.dtype.newbyteorder("=")
56 if array.dtype != native_dtype:
57 array = array.astype(native_dtype)
58 return array
61###############################################################################
62# Cache file utilities
63def _detect_compressor(fileobj):
64 """Return the compressor matching fileobj.
66 Parameters
67 ----------
68 fileobj: file object
70 Returns
71 -------
72 str in {'zlib', 'gzip', 'bz2', 'lzma', 'xz', 'compat', 'not-compressed'}
73 """
74 # Read the magic number in the first bytes of the file.
75 max_prefix_len = _get_prefixes_max_len()
76 if hasattr(fileobj, "peek"):
77 # Peek allows to read those bytes without moving the cursor in the
78 # file which.
79 first_bytes = fileobj.peek(max_prefix_len)
80 else:
81 # Fallback to seek if the fileobject is not peekable.
82 first_bytes = fileobj.read(max_prefix_len)
83 fileobj.seek(0)
85 if first_bytes.startswith(_ZFILE_PREFIX):
86 return "compat"
87 else:
88 for name, compressor in _COMPRESSORS.items():
89 if first_bytes.startswith(compressor.prefix):
90 return name
92 return "not-compressed"
95def _buffered_read_file(fobj):
96 """Return a buffered version of a read file object."""
97 return io.BufferedReader(fobj, buffer_size=_IO_BUFFER_SIZE)
100def _buffered_write_file(fobj):
101 """Return a buffered version of a write file object."""
102 return io.BufferedWriter(fobj, buffer_size=_IO_BUFFER_SIZE)
105@contextlib.contextmanager
106def _validate_fileobject_and_memmap(fileobj, filename, mmap_mode=None):
107 """Utility function opening the right fileobject from a filename.
109 The magic number is used to choose between the type of file object to open:
110 * regular file object (default)
111 * zlib file object
112 * gzip file object
113 * bz2 file object
114 * lzma file object (for xz and lzma compressor)
116 Parameters
117 ----------
118 fileobj: file object
119 filename: str
120 filename path corresponding to the fileobj parameter.
121 mmap_mode: str
122 memory map mode that should be used to open the pickle file. This
123 parameter is useful to verify that the user is not trying to one with
124 compression. Default: None.
126 Returns
127 -------
128 a tuple with a file like object, and the validated mmap_mode.
130 """
131 # Detect if the fileobj contains compressed data.
132 compressor = _detect_compressor(fileobj)
133 validated_mmap_mode = mmap_mode
135 if compressor == "compat":
136 # Compatibility with old pickle mode: simply return the input
137 # filename "as-is" and let the compatibility function be called by the
138 # caller.
139 warnings.warn(
140 "The file '%s' has been generated with a joblib "
141 "version less than 0.10. "
142 "Please regenerate this pickle file." % filename,
143 DeprecationWarning,
144 stacklevel=2,
145 )
146 yield filename, validated_mmap_mode
147 else:
148 file_must_be_closed = False
149 if compressor in _COMPRESSORS:
150 # based on the compressor detected in the file, we open the
151 # correct decompressor file object, wrapped in a buffer.
152 compressor_wrapper = _COMPRESSORS[compressor]
153 inst = compressor_wrapper.decompressor_file(fileobj)
154 fileobj = _buffered_read_file(inst)
155 file_must_be_closed = True
157 # Checking if incompatible load parameters with the type of file:
158 # mmap_mode cannot be used with compressed file or in memory buffers
159 # such as io.BytesIO.
160 if mmap_mode is not None:
161 validated_mmap_mode = None
162 if isinstance(fileobj, io.BytesIO):
163 warnings.warn(
164 "In memory persistence is not compatible with "
165 'mmap_mode "%(mmap_mode)s" flag passed. '
166 "mmap_mode option will be ignored." % locals(),
167 stacklevel=2,
168 )
169 elif compressor != "not-compressed":
170 warnings.warn(
171 'mmap_mode "%(mmap_mode)s" is not compatible '
172 "with compressed file %(filename)s. "
173 '"%(mmap_mode)s" flag will be ignored.' % locals(),
174 stacklevel=2,
175 )
176 elif not _is_raw_file(fileobj):
177 warnings.warn(
178 '"%(fileobj)r" is not a raw file, mmap_mode '
179 '"%(mmap_mode)s" flag will be ignored.' % locals(),
180 stacklevel=2,
181 )
182 else:
183 validated_mmap_mode = mmap_mode
185 try:
186 yield fileobj, validated_mmap_mode
187 finally:
188 if file_must_be_closed:
189 fileobj.close()
192def _write_fileobject(filename, compress=("zlib", 3)):
193 """Return the right compressor file object in write mode."""
194 compressmethod = compress[0]
195 compresslevel = compress[1]
197 if compressmethod in _COMPRESSORS.keys():
198 file_instance = _COMPRESSORS[compressmethod].compressor_file(
199 filename, compresslevel=compresslevel
200 )
201 return _buffered_write_file(file_instance)
202 else:
203 file_instance = _COMPRESSORS["zlib"].compressor_file(
204 filename, compresslevel=compresslevel
205 )
206 return _buffered_write_file(file_instance)
209# Utility functions/variables from numpy required for writing arrays.
210# We need at least the functions introduced in version 1.9 of numpy. Here,
211# we use the ones from numpy 1.10.2.
212BUFFER_SIZE = 2**18 # size of buffer for reading npz files in bytes
215def _read_bytes(fp, size, error_template="ran out of data"):
216 """Read from file-like object until size bytes are read.
218 TODO python2_drop: is it still needed? The docstring mentions python 2.6
219 and it looks like this can be at least simplified ...
221 Raises ValueError if not EOF is encountered before size bytes are read.
222 Non-blocking objects only supported if they derive from io objects.
224 Required as e.g. ZipExtFile in python 2.6 can return less data than
225 requested.
227 This function was taken from numpy/lib/format.py in version 1.10.2.
229 Parameters
230 ----------
231 fp: file-like object
232 size: int
233 error_template: str
235 Returns
236 -------
237 a bytes object
238 The data read in bytes.
240 """
241 data = bytes()
242 while True:
243 # io files (default in python3) return None or raise on
244 # would-block, python2 file will truncate, probably nothing can be
245 # done about that. note that regular files can't be non-blocking
246 try:
247 r = fp.read(size - len(data))
248 data += r
249 if len(r) == 0 or len(data) == size:
250 break
251 except io.BlockingIOError:
252 pass
253 if len(data) != size:
254 msg = "EOF: reading %s, expected %d bytes got %d"
255 raise ValueError(msg % (error_template, size, len(data)))
256 else:
257 return data
260def _reconstruct(*args, **kwargs):
261 # Wrapper for numpy._core.multiarray._reconstruct with backward compat
262 # for numpy 1.X
263 #
264 # XXX: Remove this function when numpy 1.X is not supported anymore
266 np_major_version = np.__version__[:2]
267 if np_major_version == "1.":
268 from numpy.core.multiarray import _reconstruct as np_reconstruct
269 elif np_major_version == "2.":
270 from numpy._core.multiarray import _reconstruct as np_reconstruct
272 return np_reconstruct(*args, **kwargs)