Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/joblib/numpy_pickle.py: 54%

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

242 statements  

1"""Utilities for fast persistence of big data, with optional compression.""" 

2 

3# Author: Gael Varoquaux <gael dot varoquaux at normalesup dot org> 

4# Copyright (c) 2009 Gael Varoquaux 

5# License: BSD Style, 3 clauses. 

6 

7import io 

8import os 

9import pickle 

10import warnings 

11 

12from .backports import make_memmap 

13from .compressor import ( 

14 _COMPRESSORS, 

15 LZ4_NOT_INSTALLED_ERROR, 

16 BinaryZlibFile, 

17 BZ2CompressorWrapper, 

18 GzipCompressorWrapper, 

19 LZ4CompressorWrapper, 

20 LZMACompressorWrapper, 

21 XZCompressorWrapper, 

22 ZlibCompressorWrapper, 

23 lz4, 

24 register_compressor, 

25) 

26 

27# For compatibility with old versions of joblib, we need ZNDArrayWrapper 

28# to be visible in the current namespace. 

29from .numpy_pickle_compat import ( 

30 NDArrayWrapper, 

31 ZNDArrayWrapper, # noqa: F401 

32 load_compatibility, 

33 readPython2Error, 

34) 

35from .numpy_pickle_utils import ( 

36 BUFFER_SIZE, 

37 Pickler, 

38 Unpickler, 

39 _ensure_native_byte_order, 

40 _read_bytes, 

41 _reconstruct, 

42 _validate_fileobject_and_memmap, 

43 _write_fileobject, 

44) 

45 

46# Register supported compressors 

47register_compressor("zlib", ZlibCompressorWrapper()) 

48register_compressor("gzip", GzipCompressorWrapper()) 

49register_compressor("bz2", BZ2CompressorWrapper()) 

50register_compressor("lzma", LZMACompressorWrapper()) 

51register_compressor("xz", XZCompressorWrapper()) 

52register_compressor("lz4", LZ4CompressorWrapper()) 

53 

54 

55############################################################################### 

56# Utility objects for persistence. 

57 

58# For convenience, 16 bytes are used to be sure to cover all the possible 

59# dtypes' alignments. For reference, see: 

60# https://numpy.org/devdocs/dev/alignment.html 

61NUMPY_ARRAY_ALIGNMENT_BYTES = 16 

62 

63 

64class NumpyArrayWrapper(object): 

65 """An object to be persisted instead of numpy arrays. 

66 

67 This object is used to hack into the pickle machinery and read numpy 

68 array data from our custom persistence format. 

69 More precisely, this object is used for: 

70 * carrying the information of the persisted array: subclass, shape, order, 

71 dtype. Those ndarray metadata are used to correctly reconstruct the array 

72 with low level numpy functions. 

73 * determining if memmap is allowed on the array. 

74 * reading the array bytes from a file. 

75 * reading the array using memorymap from a file. 

76 * writing the array bytes to a file. 

77 

78 Attributes 

79 ---------- 

80 subclass: numpy.ndarray subclass 

81 Determine the subclass of the wrapped array. 

82 shape: numpy.ndarray shape 

83 Determine the shape of the wrapped array. 

84 order: {'C', 'F'} 

85 Determine the order of wrapped array data. 'C' is for C order, 'F' is 

86 for fortran order. 

87 dtype: numpy.ndarray dtype 

88 Determine the data type of the wrapped array. 

89 allow_mmap: bool 

90 Determine if memory mapping is allowed on the wrapped array. 

91 Default: False. 

92 """ 

93 

94 def __init__( 

95 self, 

96 subclass, 

97 shape, 

98 order, 

99 dtype, 

100 allow_mmap=False, 

101 numpy_array_alignment_bytes=NUMPY_ARRAY_ALIGNMENT_BYTES, 

102 ): 

103 """Constructor. Store the useful information for later.""" 

104 self.subclass = subclass 

105 self.shape = shape 

106 self.order = order 

107 self.dtype = dtype 

108 self.allow_mmap = allow_mmap 

109 # We make numpy_array_alignment_bytes an instance attribute to allow us 

110 # to change our mind about the default alignment and still load the old 

111 # pickles (with the previous alignment) correctly 

112 self.numpy_array_alignment_bytes = numpy_array_alignment_bytes 

113 

114 def safe_get_numpy_array_alignment_bytes(self): 

115 # NumpyArrayWrapper instances loaded from joblib <= 1.1 pickles don't 

116 # have an numpy_array_alignment_bytes attribute 

117 return getattr(self, "numpy_array_alignment_bytes", None) 

118 

119 def write_array(self, array, pickler): 

120 """Write array bytes to pickler file handle. 

121 

122 This function is an adaptation of the numpy write_array function 

123 available in version 1.10.1 in numpy/lib/format.py. 

124 """ 

125 # Set buffer size to 16 MiB to hide the Python loop overhead. 

126 buffersize = max(16 * 1024**2 // array.itemsize, 1) 

127 if array.dtype.hasobject: 

128 # We contain Python objects so we cannot write out the data 

129 # directly. Instead, we will pickle it out with version 5 of the 

130 # pickle protocol. 

131 pickle.dump(array, pickler.file_handle, protocol=5) 

132 else: 

133 numpy_array_alignment_bytes = self.safe_get_numpy_array_alignment_bytes() 

134 if numpy_array_alignment_bytes is not None: 

135 current_pos = pickler.file_handle.tell() 

136 pos_after_padding_byte = current_pos + 1 

137 padding_length = numpy_array_alignment_bytes - ( 

138 pos_after_padding_byte % numpy_array_alignment_bytes 

139 ) 

140 # A single byte is written that contains the padding length in 

141 # bytes 

142 padding_length_byte = int.to_bytes( 

143 padding_length, length=1, byteorder="little" 

144 ) 

145 pickler.file_handle.write(padding_length_byte) 

146 

147 if padding_length != 0: 

148 padding = b"\xff" * padding_length 

149 pickler.file_handle.write(padding) 

150 

151 for chunk in pickler.np.nditer( 

152 array, 

153 flags=["external_loop", "buffered", "zerosize_ok"], 

154 buffersize=buffersize, 

155 order=self.order, 

156 ): 

157 pickler.file_handle.write(chunk.tobytes("C")) 

158 

159 def read_array(self, unpickler, ensure_native_byte_order): 

160 """Read array from unpickler file handle. 

161 

162 This function is an adaptation of the numpy read_array function 

163 available in version 1.10.1 in numpy/lib/format.py. 

164 """ 

165 if len(self.shape) == 0: 

166 count = 1 

167 else: 

168 # joblib issue #859: we cast the elements of self.shape to int64 to 

169 # prevent a potential overflow when computing their product. 

170 shape_int64 = [unpickler.np.int64(x) for x in self.shape] 

171 count = unpickler.np.multiply.reduce(shape_int64) 

172 # Now read the actual data. 

173 if self.dtype.hasobject: 

174 # The array contained Python objects, serialized as a nested 

175 # pickle stream. We read it with a fresh Unpickler (so the nested 

176 # stream gets its own opcode stack) but route ``find_class`` 

177 # through the *outer* unpickler instance to propagate security 

178 # harden behavior. 

179 # (see https://docs.python.org/3/library/pickle.html#pickle-restrict). 

180 inner_unpickler = Unpickler(unpickler.file_handle) 

181 inner_unpickler.find_class = unpickler.find_class 

182 array = inner_unpickler.load() 

183 else: 

184 numpy_array_alignment_bytes = self.safe_get_numpy_array_alignment_bytes() 

185 if numpy_array_alignment_bytes is not None: 

186 padding_byte = unpickler.file_handle.read(1) 

187 padding_length = int.from_bytes(padding_byte, byteorder="little") 

188 if padding_length != 0: 

189 unpickler.file_handle.read(padding_length) 

190 

191 # This is not a real file. We have to read it the 

192 # memory-intensive way. 

193 # crc32 module fails on reads greater than 2 ** 32 bytes, 

194 # breaking large reads from gzip streams. Chunk reads to 

195 # BUFFER_SIZE bytes to avoid issue and reduce memory overhead 

196 # of the read. In non-chunked case count < max_read_count, so 

197 # only one read is performed. 

198 max_read_count = BUFFER_SIZE // min(BUFFER_SIZE, self.dtype.itemsize) 

199 

200 array = unpickler.np.empty(count, dtype=self.dtype) 

201 for i in range(0, count, max_read_count): 

202 read_count = min(max_read_count, count - i) 

203 read_size = int(read_count * self.dtype.itemsize) 

204 data = _read_bytes(unpickler.file_handle, read_size, "array data") 

205 array[i : i + read_count] = unpickler.np.frombuffer( 

206 data, dtype=self.dtype, count=read_count 

207 ) 

208 del data 

209 

210 if self.order == "F": 

211 array = array.reshape(self.shape[::-1]) 

212 array = array.transpose() 

213 else: 

214 array = array.reshape(self.shape) 

215 

216 if ensure_native_byte_order: 

217 # Detect byte order mismatch and swap as needed. 

218 array = _ensure_native_byte_order(array) 

219 

220 return array 

221 

222 def read_mmap(self, unpickler): 

223 """Read an array using numpy memmap.""" 

224 current_pos = unpickler.file_handle.tell() 

225 offset = current_pos 

226 numpy_array_alignment_bytes = self.safe_get_numpy_array_alignment_bytes() 

227 

228 if numpy_array_alignment_bytes is not None: 

229 padding_byte = unpickler.file_handle.read(1) 

230 padding_length = int.from_bytes(padding_byte, byteorder="little") 

231 # + 1 is for the padding byte 

232 offset += padding_length + 1 

233 

234 if unpickler.mmap_mode == "w+": 

235 unpickler.mmap_mode = "r+" 

236 

237 marray = make_memmap( 

238 unpickler.filename, 

239 dtype=self.dtype, 

240 shape=self.shape, 

241 order=self.order, 

242 mode=unpickler.mmap_mode, 

243 offset=offset, 

244 ) 

245 # update the offset so that it corresponds to the end of the read array 

246 unpickler.file_handle.seek(offset + marray.nbytes) 

247 

248 if ( 

249 numpy_array_alignment_bytes is None 

250 and current_pos % NUMPY_ARRAY_ALIGNMENT_BYTES != 0 

251 ): 

252 message = ( 

253 f"The memmapped array {marray} loaded from the file " 

254 f"{unpickler.file_handle.name} is not byte aligned. " 

255 "This may cause segmentation faults if this memmapped array " 

256 "is used in some libraries like BLAS or PyTorch. " 

257 "To get rid of this warning, regenerate your pickle file " 

258 "with joblib >= 1.2.0. " 

259 "See https://github.com/joblib/joblib/issues/563 " 

260 "for more details" 

261 ) 

262 warnings.warn(message) 

263 

264 return marray 

265 

266 def read(self, unpickler, ensure_native_byte_order): 

267 """Read the array corresponding to this wrapper. 

268 

269 Use the unpickler to get all information to correctly read the array. 

270 

271 Parameters 

272 ---------- 

273 unpickler: NumpyUnpickler 

274 ensure_native_byte_order: bool 

275 If true, coerce the array to use the native endianness of the 

276 host system. 

277 

278 Returns 

279 ------- 

280 array: numpy.ndarray 

281 

282 """ 

283 # When requested, only use memmap mode if allowed. 

284 if unpickler.mmap_mode is not None and self.allow_mmap: 

285 assert not ensure_native_byte_order, ( 

286 "Memmaps cannot be coerced to a given byte order, " 

287 "this code path is impossible." 

288 ) 

289 array = self.read_mmap(unpickler) 

290 else: 

291 array = self.read_array(unpickler, ensure_native_byte_order) 

292 

293 # Manage array subclass case 

294 if hasattr(array, "__array_prepare__") and self.subclass not in ( 

295 unpickler.np.ndarray, 

296 unpickler.np.memmap, 

297 ): 

298 # We need to reconstruct another subclass 

299 new_array = _reconstruct(self.subclass, (0,), "b") 

300 return new_array.__array_prepare__(array) 

301 else: 

302 return array 

303 

304 

305############################################################################### 

306# Pickler classes 

307 

308 

309class NumpyPickler(Pickler): 

310 """A pickler to persist big data efficiently. 

311 

312 The main features of this object are: 

313 * persistence of numpy arrays in a single file. 

314 * optional compression with a special care on avoiding memory copies. 

315 

316 Attributes 

317 ---------- 

318 fp: file 

319 File object handle used for serializing the input object. 

320 protocol: int, optional 

321 Pickle protocol used. Default is pickle.DEFAULT_PROTOCOL. 

322 """ 

323 

324 dispatch = Pickler.dispatch.copy() 

325 

326 def __init__(self, fp, protocol=None): 

327 self.file_handle = fp 

328 self.buffered = isinstance(self.file_handle, BinaryZlibFile) 

329 

330 # By default we want a pickle protocol that only changes with 

331 # the major python version and not the minor one 

332 if protocol is None: 

333 protocol = pickle.DEFAULT_PROTOCOL 

334 

335 Pickler.__init__(self, self.file_handle, protocol=protocol) 

336 # delayed import of numpy, to avoid tight coupling 

337 try: 

338 import numpy as np 

339 except ImportError: 

340 np = None 

341 self.np = np 

342 

343 def _create_array_wrapper(self, array): 

344 """Create and returns a numpy array wrapper from a numpy array.""" 

345 order = ( 

346 "F" if (array.flags.f_contiguous and not array.flags.c_contiguous) else "C" 

347 ) 

348 allow_mmap = not self.buffered and not array.dtype.hasobject 

349 

350 kwargs = {} 

351 try: 

352 self.file_handle.tell() 

353 except io.UnsupportedOperation: 

354 kwargs = {"numpy_array_alignment_bytes": None} 

355 

356 wrapper = NumpyArrayWrapper( 

357 type(array), 

358 array.shape, 

359 order, 

360 array.dtype, 

361 allow_mmap=allow_mmap, 

362 **kwargs, 

363 ) 

364 

365 return wrapper 

366 

367 def save(self, obj): 

368 """Subclass the Pickler `save` method. 

369 

370 This is a total abuse of the Pickler class in order to use the numpy 

371 persistence function `save` instead of the default pickle 

372 implementation. The numpy array is replaced by a custom wrapper in the 

373 pickle persistence stack and the serialized array is written right 

374 after in the file. Warning: the file produced does not follow the 

375 pickle format. As such it can not be read with `pickle.load`. 

376 """ 

377 if self.np is not None and type(obj) in ( 

378 self.np.ndarray, 

379 self.np.matrix, 

380 self.np.memmap, 

381 ): 

382 if type(obj) is self.np.memmap: 

383 # Pickling doesn't work with memmapped arrays 

384 obj = self.np.asanyarray(obj) 

385 

386 # The array wrapper is pickled instead of the real array. 

387 wrapper = self._create_array_wrapper(obj) 

388 Pickler.save(self, wrapper) 

389 

390 # A framer was introduced with pickle protocol 4 and we want to 

391 # ensure the wrapper object is written before the numpy array 

392 # buffer in the pickle file. 

393 # See https://www.python.org/dev/peps/pep-3154/#framing to get 

394 # more information on the framer behavior. 

395 if self.proto >= 4: 

396 self.framer.commit_frame(force=True) 

397 

398 # And then array bytes are written right after the wrapper. 

399 wrapper.write_array(obj, self) 

400 return 

401 

402 return Pickler.save(self, obj) 

403 

404 

405class NumpyUnpickler(Unpickler): 

406 """A subclass of the Unpickler to unpickle our numpy pickles. 

407 

408 Attributes 

409 ---------- 

410 mmap_mode: str 

411 The memorymap mode to use for reading numpy arrays. 

412 file_handle: file_like 

413 File object to unpickle from. 

414 ensure_native_byte_order: bool 

415 If True, coerce the array to use the native endianness of the 

416 host system. 

417 filename: str 

418 Name of the file to unpickle from. It should correspond to file_handle. 

419 This parameter is required when using mmap_mode. 

420 np: module 

421 Reference to numpy module if numpy is installed else None. 

422 

423 """ 

424 

425 dispatch = Unpickler.dispatch.copy() 

426 

427 def __init__(self, filename, file_handle, ensure_native_byte_order, mmap_mode=None): 

428 # The next line is for backward compatibility with pickle generated 

429 # with joblib versions less than 0.10. 

430 self._dirname = os.path.dirname(filename) 

431 

432 self.mmap_mode = mmap_mode 

433 self.file_handle = file_handle 

434 # filename is required for numpy mmap mode. 

435 self.filename = filename 

436 self.compat_mode = False 

437 self.ensure_native_byte_order = ensure_native_byte_order 

438 Unpickler.__init__(self, self.file_handle) 

439 try: 

440 import numpy as np 

441 except ImportError: 

442 np = None 

443 self.np = np 

444 

445 def load_build(self): 

446 """Called to set the state of a newly created object. 

447 

448 We capture it to replace our place-holder objects, NDArrayWrapper or 

449 NumpyArrayWrapper, by the array we are interested in. We 

450 replace them directly in the stack of pickler. 

451 NDArrayWrapper is used for backward compatibility with joblib <= 0.9. 

452 """ 

453 Unpickler.load_build(self) 

454 

455 # For backward compatibility, we support NDArrayWrapper objects. 

456 if isinstance(self.stack[-1], (NDArrayWrapper, NumpyArrayWrapper)): 

457 if self.np is None: 

458 raise ImportError( 

459 "Trying to unpickle an ndarray, but numpy didn't import correctly" 

460 ) 

461 array_wrapper = self.stack.pop() 

462 # If any NDArrayWrapper is found, we switch to compatibility mode, 

463 # this will be used to raise a DeprecationWarning to the user at 

464 # the end of the unpickling. 

465 try: 

466 if isinstance(array_wrapper, NDArrayWrapper): 

467 self.compat_mode = True 

468 _array_payload = array_wrapper.read(self) 

469 else: 

470 _array_payload = array_wrapper.read( 

471 self, self.ensure_native_byte_order 

472 ) 

473 except ValueError as exc: 

474 if str(exc) == "Array can't be memory-mapped: Python objects in dtype.": 

475 raise readPython2Error(exc) 

476 raise 

477 

478 self.stack.append(_array_payload) 

479 

480 # Be careful to register our new method. 

481 dispatch[pickle.BUILD[0]] = load_build 

482 

483 

484############################################################################### 

485# Utility functions 

486 

487 

488def dump(value, filename, compress=0, protocol=None): 

489 """Persist an arbitrary Python object into one file. 

490 

491 Read more in the :ref:`User Guide <persistence>`. 

492 

493 Parameters 

494 ---------- 

495 value: any Python object 

496 The object to store to disk. 

497 filename: str, os.PathLike, or file object. 

498 The file object or path of the file in which it is to be stored. 

499 The compression method corresponding to one of the supported filename 

500 extensions ('.z', '.gz', '.bz2', '.xz' or '.lzma') will be used 

501 automatically. 

502 compress: int from 0 to 9 or bool or 2-tuple, optional 

503 Optional compression level for the data. 0 or False is no compression. 

504 Higher value means more compression, but also slower read and 

505 write times. Using a value of 3 is often a good compromise. 

506 See the notes for more details. 

507 If compress is True, the compression level used is 3. 

508 If compress is a 2-tuple, the first element must correspond to a string 

509 between supported compressors (e.g 'zlib', 'gzip', 'bz2', 'lzma' 

510 'xz'), the second element must be an integer from 0 to 9, corresponding 

511 to the compression level. 

512 protocol: int, optional 

513 Pickle protocol, see pickle.dump documentation for more details. 

514 

515 Returns 

516 ------- 

517 filenames: list of strings 

518 The list of file names in which the data is stored. If 

519 compress is false, each array is stored in a different file. 

520 

521 See Also 

522 -------- 

523 joblib.load : corresponding loader 

524 

525 Notes 

526 ----- 

527 Memmapping on load cannot be used for compressed files. Thus 

528 using compression can significantly slow down loading. In 

529 addition, compressed files take up extra memory during 

530 dump and load. 

531 

532 """ 

533 

534 if isinstance(filename, os.PathLike): 

535 filename = os.fspath(filename) 

536 

537 is_filename = isinstance(filename, str) 

538 is_fileobj = hasattr(filename, "write") 

539 

540 compress_method = "zlib" # zlib is the default compression method. 

541 if compress is True: 

542 # By default, if compress is enabled, we want the default compress 

543 # level of the compressor. 

544 compress_level = None 

545 elif isinstance(compress, tuple): 

546 # a 2-tuple was set in compress 

547 if len(compress) != 2: 

548 raise ValueError( 

549 "Compress argument tuple should contain exactly 2 elements: " 

550 "(compress method, compress level), you passed {}".format(compress) 

551 ) 

552 compress_method, compress_level = compress 

553 elif isinstance(compress, str): 

554 compress_method = compress 

555 compress_level = None # Use default compress level 

556 compress = (compress_method, compress_level) 

557 else: 

558 compress_level = compress 

559 

560 if compress_method == "lz4" and lz4 is None: 

561 raise ValueError(LZ4_NOT_INSTALLED_ERROR) 

562 

563 if ( 

564 compress_level is not None 

565 and compress_level is not False 

566 and compress_level not in range(10) 

567 ): 

568 # Raising an error if a non valid compress level is given. 

569 raise ValueError( 

570 'Non valid compress level given: "{}". Possible values are {}.'.format( 

571 compress_level, list(range(10)) 

572 ) 

573 ) 

574 

575 if compress_method not in _COMPRESSORS: 

576 # Raising an error if an unsupported compression method is given. 

577 raise ValueError( 

578 'Non valid compression method given: "{}". Possible values are {}.'.format( 

579 compress_method, _COMPRESSORS 

580 ) 

581 ) 

582 

583 if not is_filename and not is_fileobj: 

584 # People keep inverting arguments, and the resulting error is 

585 # incomprehensible 

586 raise ValueError( 

587 "Second argument should be a filename or a file-like object, " 

588 "%s (type %s) was given." % (filename, type(filename)) 

589 ) 

590 

591 if is_filename and not isinstance(compress, tuple): 

592 # In case no explicit compression was requested using both compression 

593 # method and level in a tuple and the filename has an explicit 

594 # extension, we select the corresponding compressor. 

595 

596 # unset the variable to be sure no compression level is set afterwards. 

597 compress_method = None 

598 for name, compressor in _COMPRESSORS.items(): 

599 if filename.endswith(compressor.extension): 

600 compress_method = name 

601 

602 if compress_method in _COMPRESSORS and compress_level == 0: 

603 # we choose the default compress_level in case it was not given 

604 # as an argument (using compress). 

605 compress_level = None 

606 

607 if compress_level != 0: 

608 with _write_fileobject( 

609 filename, compress=(compress_method, compress_level) 

610 ) as f: 

611 NumpyPickler(f, protocol=protocol).dump(value) 

612 elif is_filename: 

613 with open(filename, "wb") as f: 

614 NumpyPickler(f, protocol=protocol).dump(value) 

615 else: 

616 NumpyPickler(filename, protocol=protocol).dump(value) 

617 

618 # If the target container is a file object, nothing is returned. 

619 if is_fileobj: 

620 return 

621 

622 # For compatibility, the list of created filenames (e.g with one element 

623 # after 0.10.0) is returned by default. 

624 return [filename] 

625 

626 

627def _unpickle(fobj, ensure_native_byte_order, filename="", mmap_mode=None): 

628 """Internal unpickling function.""" 

629 # We are careful to open the file handle early and keep it open to 

630 # avoid race-conditions on renames. 

631 # That said, if data is stored in companion files, which can be 

632 # the case with the old persistence format, moving the directory 

633 # will create a race when joblib tries to access the companion 

634 # files. 

635 unpickler = NumpyUnpickler( 

636 filename, fobj, ensure_native_byte_order, mmap_mode=mmap_mode 

637 ) 

638 obj = None 

639 try: 

640 obj = unpickler.load() 

641 if unpickler.compat_mode: 

642 warnings.warn( 

643 "The file '%s' has been generated with a " 

644 "joblib version less than 0.10. " 

645 "Please regenerate this pickle file." % filename, 

646 DeprecationWarning, 

647 stacklevel=3, 

648 ) 

649 except UnicodeDecodeError as exc: 

650 raise readPython2Error(exc) 

651 return obj 

652 

653 

654def load_temporary_memmap(filename, mmap_mode, unlink_on_gc_collect): 

655 from ._memmapping_reducer import JOBLIB_MMAPS, add_maybe_unlink_finalizer 

656 

657 with open(filename, "rb") as f: 

658 with _validate_fileobject_and_memmap(f, filename, mmap_mode) as ( 

659 fobj, 

660 validated_mmap_mode, 

661 ): 

662 # Memmap are used for interprocess communication, which should 

663 # keep the objects untouched. We pass `ensure_native_byte_order=False` 

664 # to remain consistent with the loading behavior of non-memmaped arrays 

665 # in workers, where the byte order is preserved. 

666 # Note that we do not implement endianness change for memmaps, as this 

667 # would result in inconsistent behavior. 

668 obj = _unpickle( 

669 fobj, 

670 ensure_native_byte_order=False, 

671 filename=filename, 

672 mmap_mode=validated_mmap_mode, 

673 ) 

674 

675 JOBLIB_MMAPS.add(obj.filename) 

676 if unlink_on_gc_collect: 

677 add_maybe_unlink_finalizer(obj) 

678 return obj 

679 

680 

681def load(filename, mmap_mode=None, ensure_native_byte_order="auto"): 

682 """Reconstruct a Python object from a file persisted with joblib.dump. 

683 

684 Read more in the :ref:`User Guide <persistence>`. 

685 

686 WARNING: joblib.load relies on the pickle module and can therefore 

687 execute arbitrary Python code. It should therefore never be used 

688 to load files from untrusted sources. 

689 

690 Parameters 

691 ---------- 

692 filename: str, os.PathLike, or file object. 

693 The file object or path of the file from which to load the object 

694 mmap_mode: {None, 'r+', 'r', 'w+', 'c'}, optional 

695 If not None, the arrays are memory-mapped from the disk. This 

696 mode has no effect for compressed files. Note that in this 

697 case the reconstructed object might no longer match exactly 

698 the originally pickled object. 

699 ensure_native_byte_order: bool, or 'auto', default=='auto' 

700 If True, ensures that the byte order of the loaded arrays matches the 

701 native byte ordering (or _endianness_) of the host system. This is not 

702 compatible with memory-mapped arrays and using non-null `mmap_mode` 

703 parameter at the same time will raise an error. The default 'auto' 

704 parameter is equivalent to True if `mmap_mode` is None, else False. 

705 

706 Returns 

707 ------- 

708 result: any Python object 

709 The object stored in the file. 

710 

711 See Also 

712 -------- 

713 joblib.dump : function to save an object 

714 

715 Notes 

716 ----- 

717 

718 This function can load numpy array files saved separately during the 

719 dump. If the mmap_mode argument is given, it is passed to np.load and 

720 arrays are loaded as memmaps. As a consequence, the reconstructed 

721 object might not match the original pickled object. Note that if the 

722 file was saved with compression, the arrays cannot be memmapped. 

723 """ 

724 if ensure_native_byte_order == "auto": 

725 ensure_native_byte_order = mmap_mode is None 

726 

727 if ensure_native_byte_order and mmap_mode is not None: 

728 raise ValueError( 

729 "Native byte ordering can only be enforced if 'mmap_mode' parameter " 

730 f"is set to None, but got 'mmap_mode={mmap_mode}' instead." 

731 ) 

732 

733 if isinstance(filename, os.PathLike): 

734 filename = os.fspath(filename) 

735 

736 if hasattr(filename, "read"): 

737 fobj = filename 

738 filename = getattr(fobj, "name", "") 

739 with _validate_fileobject_and_memmap(fobj, filename, mmap_mode) as (fobj, _): 

740 obj = _unpickle(fobj, ensure_native_byte_order=ensure_native_byte_order) 

741 else: 

742 with open(filename, "rb") as f: 

743 with _validate_fileobject_and_memmap(f, filename, mmap_mode) as ( 

744 fobj, 

745 validated_mmap_mode, 

746 ): 

747 if isinstance(fobj, str): 

748 # if the returned file object is a string, this means we 

749 # try to load a pickle file generated with an version of 

750 # Joblib so we load it with joblib compatibility function. 

751 return load_compatibility(fobj) 

752 

753 # A memory-mapped array has to be mapped with the endianness 

754 # it has been written with. Other arrays are coerced to the 

755 # native endianness of the host system. 

756 obj = _unpickle( 

757 fobj, 

758 ensure_native_byte_order=ensure_native_byte_order, 

759 filename=filename, 

760 mmap_mode=validated_mmap_mode, 

761 ) 

762 

763 return obj