Do not allow pickle.load in read_array by default

"""
===============================
Improving I/O using compressors
===============================

This example compares the compressors available in Joblib. In the example,
Zlib, LZMA and LZ4 compression only are used but Joblib also supports BZ2 and
GZip compression methods.
For each compared compression method, this example dumps and reloads a
dataset fetched from an online machine-learning database. This gives 3
information: the size on disk of the compressed data, the time spent to dump
and the time spent to reload the data from disk.
"""

import os
import os.path
import time

###############################################################################
# Get some data from real-world use cases
# ---------------------------------------
#
# First fetch the benchmark dataset from an online machine-learning database
# and load it in a pandas dataframe.

import pandas as pd

url = "https://github.com/joblib/dataset/raw/main/kddcup.data.gz"
names = ("duration, protocol_type, service, flag, src_bytes, "
         "dst_bytes, land, wrong_fragment, urgent, hot, "
         "num_failed_logins, logged_in, num_compromised, "
         "root_shell, su_attempted, num_root, "
         "num_file_creations, ").split(', ')

data = pd.read_csv(url, names=names, nrows=1e6)

###############################################################################
# Dump and load the dataset without compression
# ---------------------------------------------
#
# This gives reference values for later comparison.

from joblib import dump, load

pickle_file = './pickle_data.joblib'

###############################################################################
# Start by measuring the time spent for dumping the raw data:
start = time.time()
with open(pickle_file, 'wb') as f:
    dump(data, f)
raw_dump_duration = time.time() - start
print("Raw dump duration: %0.3fs" % raw_dump_duration)

###############################################################################
# Then measure the size of the raw dumped data on disk:
raw_file_size = os.stat(pickle_file).st_size / 1e6
print("Raw dump file size: %0.3fMB" % raw_file_size)

###############################################################################
# Finally measure the time spent for loading the raw data:
start = time.time()
with open(pickle_file, 'rb') as f:
    load(f)
raw_load_duration = time.time() - start
print("Raw load duration: %0.3fs" % raw_load_duration)

###############################################################################
# Dump and load the dataset using the Zlib compression method
# -----------------------------------------------------------
#
# The compression level is using the default value, 3, which is, in general, a
# good compromise between compression and speed.

###############################################################################
# Start by measuring the time spent for dumping of the zlib data:

start = time.time()
with open(pickle_file, 'wb') as f:
    dump(data, f, compress='zlib')
zlib_dump_duration = time.time() - start
print("Zlib dump duration: %0.3fs" % zlib_dump_duration)

###############################################################################
# Then measure the size of the zlib dump data on disk:

zlib_file_size = os.stat(pickle_file).st_size / 1e6
print("Zlib file size: %0.3fMB" % zlib_file_size)

###############################################################################
# Finally measure the time spent for loading the compressed dataset:

start = time.time()
with open(pickle_file, 'rb') as f:
    load(f)
zlib_load_duration = time.time() - start
print("Zlib load duration: %0.3fs" % zlib_load_duration)

###############################################################################
# .. note:: The compression format is detected automatically by Joblib.
#    The compression format is identified by the standard magic number present
#    at the beginning of the file. Joblib uses this information to determine
#    the compression method used.
#    This is the case for all compression methods supported by Joblib.

###############################################################################
# Dump and load the dataset using the LZMA compression method
# -----------------------------------------------------------
#
# LZMA compression method has a very good compression rate but at the cost
# of being very slow.
# In this example, a light compression level, e.g. 3, is used to speed up a
# bit the dump/load cycle.

###############################################################################
# Start by measuring the time spent for dumping the lzma data:

start = time.time()
with open(pickle_file, 'wb') as f:
    dump(data, f, compress=('lzma', 3))
lzma_dump_duration = time.time() - start
print("LZMA dump duration: %0.3fs" % lzma_dump_duration)

###############################################################################
# Then measure the size of the lzma dump data on disk:

lzma_file_size = os.stat(pickle_file).st_size / 1e6
print("LZMA file size: %0.3fMB" % lzma_file_size)

###############################################################################
# Finally measure the time spent for loading the lzma data:

start = time.time()
with open(pickle_file, 'rb') as f:
    load(f)
lzma_load_duration = time.time() - start
print("LZMA load duration: %0.3fs" % lzma_load_duration)

###############################################################################
# Dump and load the dataset using the LZ4 compression method
# ----------------------------------------------------------
#
# LZ4 compression method is known to be one of the fastest available
# compression method but with a compression rate a bit lower than Zlib. In
# most of the cases, this method is a good choice.

###############################################################################
# .. note:: In order to use LZ4 compression with Joblib, the
#    `lz4 <https://pypi.python.org/pypi/lz4>`_ package must be installed
#    on the system.

###############################################################################
# Start by measuring the time spent for dumping the lz4 data:

start = time.time()
with open(pickle_file, 'wb') as f:
    dump(data, f, compress='lz4')
lz4_dump_duration = time.time() - start
print("LZ4 dump duration: %0.3fs" % lz4_dump_duration)

###############################################################################
# Then measure the size of the lz4 dump data on disk:

lz4_file_size = os.stat(pickle_file).st_size / 1e6
print("LZ4 file size: %0.3fMB" % lz4_file_size)

###############################################################################
# Finally measure the time spent for loading the lz4 data:

start = time.time()
with open(pickle_file, 'rb') as f:
    load(f)
lz4_load_duration = time.time() - start
print("LZ4 load duration: %0.3fs" % lz4_load_duration)

###############################################################################
# Comparing the results
# ---------------------

import numpy as np
import matplotlib.pyplot as plt

N = 4
load_durations = (raw_load_duration, lz4_load_duration, zlib_load_duration,
                  lzma_load_duration)
dump_durations = (raw_dump_duration, lz4_dump_duration, zlib_dump_duration,
                  lzma_dump_duration)
file_sizes = (raw_file_size, lz4_file_size, zlib_file_size, lzma_file_size)
ind = np.arange(N)
width = 0.5

plt.figure(1, figsize=(5, 4))
p1 = plt.bar(ind, dump_durations, width)
p2 = plt.bar(ind, load_durations, width, bottom=dump_durations)
plt.ylabel('Time in seconds')
plt.title('Dump and load durations')
plt.xticks(ind, ('Raw', 'LZ4', 'Zlib', 'LZMA'))
plt.yticks(np.arange(0, lzma_load_duration + lzma_dump_duration))
plt.legend((p1[0], p2[0]), ('Dump duration', 'Load duration'))

###############################################################################
# Compared with other compressors, LZ4 is clearly the fastest, especially for
# dumping compressed data on disk. In this particular case, it can even be
# faster than the raw dump.
# Also note that dump and load durations depend on the I/O speed of the
# underlying storage: for example, with SSD hard drives the LZ4 compression
# will be slightly slower than raw dump/load, whereas with spinning hard disk
# drives (HDD) or remote storage (NFS), LZ4 is faster in general.
#
# LZMA and Zlib, even if always slower for dumping data, are quite fast when
# re-loading compressed data from disk.

plt.figure(2, figsize=(5, 4))
plt.bar(ind, file_sizes, width, log=True)
plt.ylabel('File size in MB')
plt.xticks(ind, ('Raw', 'LZ4', 'Zlib', 'LZMA'))

###############################################################################
# Compressed data obviously takes a lot less space on disk than raw data. LZMA
# is the best compression method in terms of compression rate. Zlib also has a
# better compression rate than LZ4.

plt.show()

###############################################################################
# Clear the pickle file
# ---------------------

import os
os.remove(pickle_file)

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

# Author: Gael Varoquaux <gael dot varoquaux at normalesup dot org>
# Copyright (c) 2009 Gael Varoquaux
# License: BSD Style, 3 clauses.

import pickle
import os
import warnings
import io
from pathlib import Path

from .compressor import lz4, LZ4_NOT_INSTALLED_ERROR
from .compressor import _COMPRESSORS, register_compressor, BinaryZlibFile
from .compressor import (ZlibCompressorWrapper, GzipCompressorWrapper,
                         BZ2CompressorWrapper, LZMACompressorWrapper,
                         XZCompressorWrapper, LZ4CompressorWrapper)
from .numpy_pickle_utils import Unpickler, Pickler
from .numpy_pickle_utils import _read_fileobject, _write_fileobject
from .numpy_pickle_utils import _read_bytes, BUFFER_SIZE
from .numpy_pickle_utils import _ensure_native_byte_order
from .numpy_pickle_compat import load_compatibility
from .numpy_pickle_compat import NDArrayWrapper
# For compatibility with old versions of joblib, we need ZNDArrayWrapper
# to be visible in the current namespace.
# Explicitly skipping next line from flake8 as it triggers an F401 warning
# which we don't care.
from .numpy_pickle_compat import ZNDArrayWrapper  # noqa
from .backports import make_memmap

# Register supported compressors
register_compressor('zlib', ZlibCompressorWrapper())
register_compressor('gzip', GzipCompressorWrapper())
register_compressor('bz2', BZ2CompressorWrapper())
register_compressor('lzma', LZMACompressorWrapper())
register_compressor('xz', XZCompressorWrapper())
register_compressor('lz4', LZ4CompressorWrapper())

###############################################################################
# Utility objects for persistence.

# For convenience, 16 bytes are used to be sure to cover all the possible
# dtypes' alignments. For reference, see:
# https://numpy.org/devdocs/dev/alignment.html
NUMPY_ARRAY_ALIGNMENT_BYTES = 16

class NumpyArrayWrapper(object):
    """An object to be persisted instead of numpy arrays.

    This object is used to hack into the pickle machinery and read numpy
    array data from our custom persistence format.
    More precisely, this object is used for:
    * carrying the information of the persisted array: subclass, shape, order,
    dtype. Those ndarray metadata are used to correctly reconstruct the array
    with low level numpy functions.
    * determining if memmap is allowed on the array.
    * reading the array bytes from a file.
    * reading the array using memorymap from a file.
    * writing the array bytes to a file.

    Attributes
    ----------
    subclass: numpy.ndarray subclass
        Determine the subclass of the wrapped array.
    shape: numpy.ndarray shape
        Determine the shape of the wrapped array.
    order: {'C', 'F'}
        Determine the order of wrapped array data. 'C' is for C order, 'F' is
        for fortran order.
    dtype: numpy.ndarray dtype
        Determine the data type of the wrapped array.
    allow_mmap: bool
        Determine if memory mapping is allowed on the wrapped array.
        Default: False.
    """

    def __init__(self, subclass, shape, order, dtype, allow_mmap=False,
                 numpy_array_alignment_bytes=NUMPY_ARRAY_ALIGNMENT_BYTES):
        """Constructor. Store the useful information for later."""
        self.subclass = subclass
        self.shape = shape
        self.order = order
        self.dtype = dtype
        self.allow_mmap = allow_mmap
        # We make numpy_array_alignment_bytes an instance attribute to allow us
        # to change our mind about the default alignment and still load the old
        # pickles (with the previous alignment) correctly
        self.numpy_array_alignment_bytes = numpy_array_alignment_bytes

    def safe_get_numpy_array_alignment_bytes(self):
        # NumpyArrayWrapper instances loaded from joblib <= 1.1 pickles don't
        # have an numpy_array_alignment_bytes attribute
        return getattr(self, 'numpy_array_alignment_bytes', None)

    def write_array(self, array, pickler):
        """Write array bytes to pickler file handle.

        This function is an adaptation of the numpy write_array function
        available in version 1.10.1 in numpy/lib/format.py.
        """
        # Set buffer size to 16 MiB to hide the Python loop overhead.
        buffersize = max(16 * 1024 ** 2 // array.itemsize, 1)
        if array.dtype.hasobject:
            # We contain Python objects so we cannot write out the data
            # directly. Instead, we will pickle it out with version 2 of the
            # pickle protocol.
            pickle.dump(array, pickler.file_handle, protocol=2)
        else:
            numpy_array_alignment_bytes = \
                self.safe_get_numpy_array_alignment_bytes()
            if numpy_array_alignment_bytes is not None:
                current_pos = pickler.file_handle.tell()
                pos_after_padding_byte = current_pos + 1
                padding_length = numpy_array_alignment_bytes - (
                    pos_after_padding_byte % numpy_array_alignment_bytes)
                # A single byte is written that contains the padding length in
                # bytes
                padding_length_byte = int.to_bytes(
                    padding_length, length=1, byteorder='little')
                pickler.file_handle.write(padding_length_byte)

                if padding_length != 0:
                    padding = b'\xff' * padding_length
                    pickler.file_handle.write(padding)

            for chunk in pickler.np.nditer(array,
                                           flags=['external_loop',
                                                  'buffered',
                                                  'zerosize_ok'],
                                           buffersize=buffersize,
                                           order=self.order):
                pickler.file_handle.write(chunk.tobytes('C'))

    def read_array(self, unpickler):
        """Read array from unpickler file handle.

        This function is an adaptation of the numpy read_array function
        available in version 1.10.1 in numpy/lib/format.py.
        """
        if len(self.shape) == 0:
            count = 1
        else:
            # joblib issue #859: we cast the elements of self.shape to int64 to
            # prevent a potential overflow when computing their product.
            shape_int64 = [unpickler.np.int64(x) for x in self.shape]
            count = unpickler.np.multiply.reduce(shape_int64)
        # Now read the actual data.
        if self.dtype.hasobject:
            # The array contained Python objects. We need to unpickle the data.
            array = pickle.load(unpickler.file_handle)
        else:
            numpy_array_alignment_bytes = \
                self.safe_get_numpy_array_alignment_bytes()
            if numpy_array_alignment_bytes is not None:
                padding_byte = unpickler.file_handle.read(1)
                padding_length = int.from_bytes(
                    padding_byte, byteorder='little')
                if padding_length != 0:
                    unpickler.file_handle.read(padding_length)

            # This is not a real file. We have to read it the
            # memory-intensive way.
            # crc32 module fails on reads greater than 2 ** 32 bytes,
            # breaking large reads from gzip streams. Chunk reads to
            # BUFFER_SIZE bytes to avoid issue and reduce memory overhead
            # of the read. In non-chunked case count < max_read_count, so
            # only one read is performed.
            max_read_count = BUFFER_SIZE // min(BUFFER_SIZE,
                                                self.dtype.itemsize)

            array = unpickler.np.empty(count, dtype=self.dtype)
            for i in range(0, count, max_read_count):
                read_count = min(max_read_count, count - i)
                read_size = int(read_count * self.dtype.itemsize)
                data = _read_bytes(unpickler.file_handle,
                                   read_size, "array data")
                array[i:i + read_count] = \
                    unpickler.np.frombuffer(data, dtype=self.dtype,
                                            count=read_count)
                del data

            if self.order == 'F':
                array.shape = self.shape[::-1]
                array = array.transpose()
            else:
                array.shape = self.shape

        # Detect byte order mismatch and swap as needed.
        return _ensure_native_byte_order(array)

    def read_mmap(self, unpickler):
        """Read an array using numpy memmap."""
        current_pos = unpickler.file_handle.tell()
        offset = current_pos
        numpy_array_alignment_bytes = \
            self.safe_get_numpy_array_alignment_bytes()

        if numpy_array_alignment_bytes is not None:
            padding_byte = unpickler.file_handle.read(1)
            padding_length = int.from_bytes(padding_byte, byteorder='little')
            # + 1 is for the padding byte
            offset += padding_length + 1

        if unpickler.mmap_mode == 'w+':
            unpickler.mmap_mode = 'r+'

        marray = make_memmap(unpickler.filename,
                             dtype=self.dtype,
                             shape=self.shape,
                             order=self.order,
                             mode=unpickler.mmap_mode,
                             offset=offset)
        # update the offset so that it corresponds to the end of the read array
        unpickler.file_handle.seek(offset + marray.nbytes)

        if (numpy_array_alignment_bytes is None and
                current_pos % NUMPY_ARRAY_ALIGNMENT_BYTES != 0):
            message = (
                f'The memmapped array {marray} loaded from the file '
                f'{unpickler.file_handle.name} is not byte aligned. '
                'This may cause segmentation faults if this memmapped array '
                'is used in some libraries like BLAS or PyTorch. '
                'To get rid of this warning, regenerate your pickle file '
                'with joblib >= 1.2.0. '
                'See https://github.com/joblib/joblib/issues/563 '
                'for more details'
            )
            warnings.warn(message)

        return _ensure_native_byte_order(marray)

    def read(self, unpickler):
        """Read the array corresponding to this wrapper.

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

        Parameters
        ----------
        unpickler: NumpyUnpickler

        Returns
        -------
        array: numpy.ndarray

        """
        # When requested, only use memmap mode if allowed.
        if unpickler.mmap_mode is not None and self.allow_mmap:
            array = self.read_mmap(unpickler)
        else:
            array = self.read_array(unpickler)

        # Manage array subclass case
        if (hasattr(array, '__array_prepare__') and
            self.subclass not in (unpickler.np.ndarray,
                                  unpickler.np.memmap)):
            # We need to reconstruct another subclass
            new_array = unpickler.np.core.multiarray._reconstruct(
                self.subclass, (0,), 'b')
            return new_array.__array_prepare__(array)
        else:
            return array

###############################################################################
# Pickler classes

class NumpyPickler(Pickler):
    """A pickler to persist big data efficiently.

    The main features of this object are:
    * persistence of numpy arrays in a single file.
    * optional compression with a special care on avoiding memory copies.

    Attributes
    ----------
    fp: file
        File object handle used for serializing the input object.
    protocol: int, optional
        Pickle protocol used. Default is pickle.DEFAULT_PROTOCOL.
    """

    dispatch = Pickler.dispatch.copy()

    def __init__(self, fp, protocol=None):
        self.file_handle = fp
        self.buffered = isinstance(self.file_handle, BinaryZlibFile)

        # By default we want a pickle protocol that only changes with
        # the major python version and not the minor one
        if protocol is None:
            protocol = pickle.DEFAULT_PROTOCOL

        Pickler.__init__(self, self.file_handle, protocol=protocol)
        # delayed import of numpy, to avoid tight coupling
        try:
            import numpy as np
        except ImportError:
            np = None
        self.np = np

    def _create_array_wrapper(self, array):
        """Create and returns a numpy array wrapper from a numpy array."""
        order = 'F' if (array.flags.f_contiguous and
                        not array.flags.c_contiguous) else 'C'
        allow_mmap = not self.buffered and not array.dtype.hasobject

        kwargs = {}
        try:
            self.file_handle.tell()
        except io.UnsupportedOperation:
            kwargs = {'numpy_array_alignment_bytes': None}

        wrapper = NumpyArrayWrapper(type(array),
                                    array.shape, order, array.dtype,
                                    allow_mmap=allow_mmap,
                                    **kwargs)

        return wrapper

    def save(self, obj):
        """Subclass the Pickler `save` method.

        This is a total abuse of the Pickler class in order to use the numpy
        persistence function `save` instead of the default pickle
        implementation. The numpy array is replaced by a custom wrapper in the
        pickle persistence stack and the serialized array is written right
        after in the file. Warning: the file produced does not follow the
        pickle format. As such it can not be read with `pickle.load`.
        """
        if self.np is not None and type(obj) in (self.np.ndarray,
                                                 self.np.matrix,
                                                 self.np.memmap):
            if type(obj) is self.np.memmap:
                # Pickling doesn't work with memmapped arrays
                obj = self.np.asanyarray(obj)

            # The array wrapper is pickled instead of the real array.
            wrapper = self._create_array_wrapper(obj)
            Pickler.save(self, wrapper)

            # A framer was introduced with pickle protocol 4 and we want to
            # ensure the wrapper object is written before the numpy array
            # buffer in the pickle file.
            # See https://www.python.org/dev/peps/pep-3154/#framing to get
            # more information on the framer behavior.
            if self.proto >= 4:
                self.framer.commit_frame(force=True)

            # And then array bytes are written right after the wrapper.
            wrapper.write_array(obj, self)
            return

        return Pickler.save(self, obj)

class NumpyUnpickler(Unpickler):
    """A subclass of the Unpickler to unpickle our numpy pickles.

    Attributes
    ----------
    mmap_mode: str
        The memorymap mode to use for reading numpy arrays.
    file_handle: file_like
        File object to unpickle from.
    filename: str
        Name of the file to unpickle from. It should correspond to file_handle.
        This parameter is required when using mmap_mode.
    np: module
        Reference to numpy module if numpy is installed else None.

    """

    dispatch = Unpickler.dispatch.copy()

    def __init__(self, filename, file_handle, mmap_mode=None):
        # The next line is for backward compatibility with pickle generated
        # with joblib versions less than 0.10.
        self._dirname = os.path.dirname(filename)

        self.mmap_mode = mmap_mode
        self.file_handle = file_handle
        # filename is required for numpy mmap mode.
        self.filename = filename
        self.compat_mode = False
        Unpickler.__init__(self, self.file_handle)
        try:
            import numpy as np
        except ImportError:
            np = None
        self.np = np

    def load_build(self):
        """Called to set the state of a newly created object.

        We capture it to replace our place-holder objects, NDArrayWrapper or
        NumpyArrayWrapper, by the array we are interested in. We
        replace them directly in the stack of pickler.
        NDArrayWrapper is used for backward compatibility with joblib <= 0.9.
        """
        Unpickler.load_build(self)

        # For backward compatibility, we support NDArrayWrapper objects.
        if isinstance(self.stack[-1], (NDArrayWrapper, NumpyArrayWrapper)):
            if self.np is None:
                raise ImportError("Trying to unpickle an ndarray, "
                                  "but numpy didn't import correctly")
            array_wrapper = self.stack.pop()
            # If any NDArrayWrapper is found, we switch to compatibility mode,
            # this will be used to raise a DeprecationWarning to the user at
            # the end of the unpickling.
            if isinstance(array_wrapper, NDArrayWrapper):
                self.compat_mode = True
            self.stack.append(array_wrapper.read(self))

    # Be careful to register our new method.
    dispatch[pickle.BUILD[0]] = load_build

###############################################################################
# Utility functions

def dump(value, filename, compress=0, protocol=None, cache_size=None):
    """Persist an arbitrary Python object into one file.

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

    Parameters
    ----------
    value: any Python object
        The object to store to disk.
    filename: str, pathlib.Path, or file object.
        The file object or path of the file in which it is to be stored.
        The compression method corresponding to one of the supported filename
        extensions ('.z', '.gz', '.bz2', '.xz' or '.lzma') will be used
        automatically.
    compress: int from 0 to 9 or bool or 2-tuple, optional
        Optional compression level for the data. 0 or False is no compression.
        Higher value means more compression, but also slower read and
        write times. Using a value of 3 is often a good compromise.
        See the notes for more details.
        If compress is True, the compression level used is 3.
        If compress is a 2-tuple, the first element must correspond to a string
        between supported compressors (e.g 'zlib', 'gzip', 'bz2', 'lzma'
        'xz'), the second element must be an integer from 0 to 9, corresponding
        to the compression level.
    protocol: int, optional
        Pickle protocol, see pickle.dump documentation for more details.
    cache_size: positive int, optional
        This option is deprecated in 0.10 and has no effect.

    Returns
    -------
    filenames: list of strings
        The list of file names in which the data is stored. If
        compress is false, each array is stored in a different file.

    See Also
    --------
    joblib.load : corresponding loader

    Notes
    -----
    Memmapping on load cannot be used for compressed files. Thus
    using compression can significantly slow down loading. In
    addition, compressed files take up extra memory during
    dump and load.

    """

    if Path is not None and isinstance(filename, Path):
        filename = str(filename)

    is_filename = isinstance(filename, str)
    is_fileobj = hasattr(filename, "write")

    compress_method = 'zlib'  # zlib is the default compression method.
    if compress is True:
        # By default, if compress is enabled, we want the default compress
        # level of the compressor.
        compress_level = None
    elif isinstance(compress, tuple):
        # a 2-tuple was set in compress
        if len(compress) != 2:
            raise ValueError(
                'Compress argument tuple should contain exactly 2 elements: '
                '(compress method, compress level), you passed {}'
                .format(compress))
        compress_method, compress_level = compress
    elif isinstance(compress, str):
        compress_method = compress
        compress_level = None  # Use default compress level
        compress = (compress_method, compress_level)
    else:
        compress_level = compress

    if compress_method == 'lz4' and lz4 is None:
        raise ValueError(LZ4_NOT_INSTALLED_ERROR)

    if (compress_level is not None and
            compress_level is not False and
            compress_level not in range(10)):
        # Raising an error if a non valid compress level is given.
        raise ValueError(
            'Non valid compress level given: "{}". Possible values are '
            '{}.'.format(compress_level, list(range(10))))

    if compress_method not in _COMPRESSORS:
        # Raising an error if an unsupported compression method is given.
        raise ValueError(
            'Non valid compression method given: "{}". Possible values are '
            '{}.'.format(compress_method, _COMPRESSORS))

    if not is_filename and not is_fileobj:
        # People keep inverting arguments, and the resulting error is
        # incomprehensible
        raise ValueError(
            'Second argument should be a filename or a file-like object, '
            '%s (type %s) was given.'
            % (filename, type(filename))
        )

    if is_filename and not isinstance(compress, tuple):
        # In case no explicit compression was requested using both compression
        # method and level in a tuple and the filename has an explicit
        # extension, we select the corresponding compressor.

        # unset the variable to be sure no compression level is set afterwards.
        compress_method = None
        for name, compressor in _COMPRESSORS.items():
            if filename.endswith(compressor.extension):
                compress_method = name

        if compress_method in _COMPRESSORS and compress_level == 0:
            # we choose the default compress_level in case it was not given
            # as an argument (using compress).
            compress_level = None

    if cache_size is not None:
        # Cache size is deprecated starting from version 0.10
        warnings.warn("Please do not set 'cache_size' in joblib.dump, "
                      "this parameter has no effect and will be removed. "
                      "You used 'cache_size={}'".format(cache_size),
                      DeprecationWarning, stacklevel=2)

    if compress_level != 0:
        with _write_fileobject(filename, compress=(compress_method,
                                                   compress_level)) as f:
            NumpyPickler(f, protocol=protocol).dump(value)
    elif is_filename:
        with open(filename, 'wb') as f:
            NumpyPickler(f, protocol=protocol).dump(value)
    else:
        NumpyPickler(filename, protocol=protocol).dump(value)

    # If the target container is a file object, nothing is returned.
    if is_fileobj:
        return

    # For compatibility, the list of created filenames (e.g with one element
    # after 0.10.0) is returned by default.
    return [filename]

def _unpickle(fobj, filename="", mmap_mode=None):
    """Internal unpickling function."""
    # We are careful to open the file handle early and keep it open to
    # avoid race-conditions on renames.
    # That said, if data is stored in companion files, which can be
    # the case with the old persistence format, moving the directory
    # will create a race when joblib tries to access the companion
    # files.
    unpickler = NumpyUnpickler(filename, fobj, mmap_mode=mmap_mode)
    obj = None
    try:
        obj = unpickler.load()
        if unpickler.compat_mode:
            warnings.warn("The file '%s' has been generated with a "
                          "joblib version less than 0.10. "
                          "Please regenerate this pickle file."
                          % filename,
                          DeprecationWarning, stacklevel=3)
    except UnicodeDecodeError as exc:
        # More user-friendly error message
        new_exc = ValueError(
            'You may be trying to read with '
            'python 3 a joblib pickle generated with python 2. '
            'This feature is not supported by joblib.')
        new_exc.__cause__ = exc
        raise new_exc
    return obj

def load_temporary_memmap(filename, mmap_mode, unlink_on_gc_collect):
    from ._memmapping_reducer import JOBLIB_MMAPS, add_maybe_unlink_finalizer
    obj = load(filename, mmap_mode)
    JOBLIB_MMAPS.add(obj.filename)
    if unlink_on_gc_collect:
        add_maybe_unlink_finalizer(obj)
    return obj

def load(filename, mmap_mode=None):
    """Reconstruct a Python object from a file persisted with joblib.dump.

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

    WARNING: joblib.load relies on the pickle module and can therefore
    execute arbitrary Python code. It should therefore never be used
    to load files from untrusted sources.

    Parameters
    ----------
    filename: str, pathlib.Path, or file object.
        The file object or path of the file from which to load the object
    mmap_mode: {None, 'r+', 'r', 'w+', 'c'}, optional
        If not None, the arrays are memory-mapped from the disk. This
        mode has no effect for compressed files. Note that in this
        case the reconstructed object might no longer match exactly
        the originally pickled object.

    Returns
    -------
    result: any Python object
        The object stored in the file.

    See Also
    --------
    joblib.dump : function to save an object

    Notes
    -----

    This function can load numpy array files saved separately during the
    dump. If the mmap_mode argument is given, it is passed to np.load and
    arrays are loaded as memmaps. As a consequence, the reconstructed
    object might not match the original pickled object. Note that if the
    file was saved with compression, the arrays cannot be memmapped.
    """
    if Path is not None and isinstance(filename, Path):
        filename = str(filename)

    if hasattr(filename, "read"):
        fobj = filename
        filename = getattr(fobj, 'name', '')
        with _read_fileobject(fobj, filename, mmap_mode) as fobj:
            obj = _unpickle(fobj)
    else:
        with open(filename, 'rb') as f:
            with _read_fileobject(f, filename, mmap_mode) as fobj:
                if isinstance(fobj, str):
                    # if the returned file object is a string, this means we
                    # try to load a pickle file generated with an version of
                    # Joblib so we load it with joblib compatibility function.
                    return load_compatibility(fobj)

                obj = _unpickle(fobj, filename, mmap_mode)
    return obj

"""Test the numpy pickler as a replacement of the standard pickler."""

import copy
import os
import random
import re
import io
import sys
import warnings
import gzip
import zlib
import bz2
import pickle
import socket
from contextlib import closing
import mmap
from pathlib import Path

try:
    import lzma
except ImportError:
    lzma = None

import pytest

from joblib.test.common import np, with_numpy, with_lz4, without_lz4
from joblib.test.common import with_memory_profiler, memory_used
from joblib.testing import parametrize, raises, warns

# numpy_pickle is not a drop-in replacement of pickle, as it takes
# filenames instead of open files as arguments.
from joblib import numpy_pickle, register_compressor
from joblib.test import data

from joblib.numpy_pickle_utils import _IO_BUFFER_SIZE
from joblib.numpy_pickle_utils import _detect_compressor
from joblib.numpy_pickle_utils import _is_numpy_array_byte_order_mismatch
from joblib.numpy_pickle_utils import _ensure_native_byte_order
from joblib.compressor import (_COMPRESSORS, _LZ4_PREFIX, CompressorWrapper,
                               LZ4_NOT_INSTALLED_ERROR, BinaryZlibFile)

###############################################################################
# Define a list of standard types.
# Borrowed from dill, initial author: Micheal McKerns:
# http://dev.danse.us/trac/pathos/browser/dill/dill_test2.py

typelist = []

# testing types
_none = None
typelist.append(_none)
_type = type
typelist.append(_type)
_bool = bool(1)
typelist.append(_bool)
_int = int(1)
typelist.append(_int)
_float = float(1)
typelist.append(_float)
_complex = complex(1)
typelist.append(_complex)
_string = str(1)
typelist.append(_string)
_tuple = ()
typelist.append(_tuple)
_list = []
typelist.append(_list)
_dict = {}
typelist.append(_dict)
_builtin = len
typelist.append(_builtin)

def _function(x):
    yield x

class _class:
    def _method(self):
        pass

class _newclass(object):
    def _method(self):
        pass

typelist.append(_function)
typelist.append(_class)
typelist.append(_newclass)  # <type 'type'>
_instance = _class()
typelist.append(_instance)
_object = _newclass()
typelist.append(_object)  # <type 'class'>

###############################################################################
# Tests

@parametrize('compress', [0, 1])
@parametrize('member', typelist)
def test_standard_types(tmpdir, compress, member):
    # Test pickling and saving with standard types.
    filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump(member, filename, compress=compress)
    _member = numpy_pickle.load(filename)
    # We compare the pickled instance to the reloaded one only if it
    # can be compared to a copied one
    if member == copy.deepcopy(member):
        assert member == _member

def test_value_error():
    # Test inverting the input arguments to dump
    with raises(ValueError):
        numpy_pickle.dump('foo', dict())

@parametrize('wrong_compress', [-1, 10, dict()])
def test_compress_level_error(wrong_compress):
    # Verify that passing an invalid compress argument raises an error.
    exception_msg = ('Non valid compress level given: '
                     '"{0}"'.format(wrong_compress))
    with raises(ValueError) as excinfo:
        numpy_pickle.dump('dummy', 'foo', compress=wrong_compress)
    excinfo.match(exception_msg)

@with_numpy
@parametrize('compress', [False, True, 0, 3, 'zlib'])
def test_numpy_persistence(tmpdir, compress):
    filename = tmpdir.join('test.pkl').strpath
    rnd = np.random.RandomState(0)
    a = rnd.random_sample((10, 2))
    # We use 'a.T' to have a non C-contiguous array.
    for index, obj in enumerate(((a,), (a.T,), (a, a), [a, a, a])):
        filenames = numpy_pickle.dump(obj, filename, compress=compress)

        # All is cached in one file
        assert len(filenames) == 1
        # Check that only one file was created
        assert filenames[0] == filename
        # Check that this file does exist
        assert os.path.exists(filenames[0])

        # Unpickle the object
        obj_ = numpy_pickle.load(filename)
        # Check that the items are indeed arrays
        for item in obj_:
            assert isinstance(item, np.ndarray)
        # And finally, check that all the values are equal.
        np.testing.assert_array_equal(np.array(obj), np.array(obj_))

    # Now test with an array subclass
    obj = np.memmap(filename + 'mmap', mode='w+', shape=4, dtype=np.float64)
    filenames = numpy_pickle.dump(obj, filename, compress=compress)
    # All is cached in one file
    assert len(filenames) == 1

    obj_ = numpy_pickle.load(filename)
    if (type(obj) is not np.memmap and
            hasattr(obj, '__array_prepare__')):
        # We don't reconstruct memmaps
        assert isinstance(obj_, type(obj))

    np.testing.assert_array_equal(obj_, obj)

    # Test with an object containing multiple numpy arrays
    obj = ComplexTestObject()
    filenames = numpy_pickle.dump(obj, filename, compress=compress)
    # All is cached in one file
    assert len(filenames) == 1

    obj_loaded = numpy_pickle.load(filename)
    assert isinstance(obj_loaded, type(obj))
    np.testing.assert_array_equal(obj_loaded.array_float, obj.array_float)
    np.testing.assert_array_equal(obj_loaded.array_int, obj.array_int)
    np.testing.assert_array_equal(obj_loaded.array_obj, obj.array_obj)

@with_numpy
def test_numpy_persistence_bufferred_array_compression(tmpdir):
    big_array = np.ones((_IO_BUFFER_SIZE + 100), dtype=np.uint8)
    filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump(big_array, filename, compress=True)
    arr_reloaded = numpy_pickle.load(filename)

    np.testing.assert_array_equal(big_array, arr_reloaded)

@with_numpy
def test_memmap_persistence(tmpdir):
    rnd = np.random.RandomState(0)
    a = rnd.random_sample(10)
    filename = tmpdir.join('test1.pkl').strpath
    numpy_pickle.dump(a, filename)
    b = numpy_pickle.load(filename, mmap_mode='r')

    assert isinstance(b, np.memmap)

    # Test with an object containing multiple numpy arrays
    filename = tmpdir.join('test2.pkl').strpath
    obj = ComplexTestObject()
    numpy_pickle.dump(obj, filename)
    obj_loaded = numpy_pickle.load(filename, mmap_mode='r')
    assert isinstance(obj_loaded, type(obj))
    assert isinstance(obj_loaded.array_float, np.memmap)
    assert not obj_loaded.array_float.flags.writeable
    assert isinstance(obj_loaded.array_int, np.memmap)
    assert not obj_loaded.array_int.flags.writeable
    # Memory map not allowed for numpy object arrays
    assert not isinstance(obj_loaded.array_obj, np.memmap)
    np.testing.assert_array_equal(obj_loaded.array_float,
                                  obj.array_float)
    np.testing.assert_array_equal(obj_loaded.array_int,
                                  obj.array_int)
    np.testing.assert_array_equal(obj_loaded.array_obj,
                                  obj.array_obj)

    # Test we can write in memmapped arrays
    obj_loaded = numpy_pickle.load(filename, mmap_mode='r+')
    assert obj_loaded.array_float.flags.writeable
    obj_loaded.array_float[0:10] = 10.0
    assert obj_loaded.array_int.flags.writeable
    obj_loaded.array_int[0:10] = 10

    obj_reloaded = numpy_pickle.load(filename, mmap_mode='r')
    np.testing.assert_array_equal(obj_reloaded.array_float,
                                  obj_loaded.array_float)
    np.testing.assert_array_equal(obj_reloaded.array_int,
                                  obj_loaded.array_int)

    # Test w+ mode is caught and the mode has switched to r+
    numpy_pickle.load(filename, mmap_mode='w+')
    assert obj_loaded.array_int.flags.writeable
    assert obj_loaded.array_int.mode == 'r+'
    assert obj_loaded.array_float.flags.writeable
    assert obj_loaded.array_float.mode == 'r+'

@with_numpy
def test_memmap_persistence_mixed_dtypes(tmpdir):
    # loading datastructures that have sub-arrays with dtype=object
    # should not prevent memmapping on fixed size dtype sub-arrays.
    rnd = np.random.RandomState(0)
    a = rnd.random_sample(10)
    b = np.array([1, 'b'], dtype=object)
    construct = (a, b)
    filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump(construct, filename)
    a_clone, b_clone = numpy_pickle.load(filename, mmap_mode='r')

    # the floating point array has been memory mapped
    assert isinstance(a_clone, np.memmap)

    # the object-dtype array has been loaded in memory
    assert not isinstance(b_clone, np.memmap)

@with_numpy
def test_masked_array_persistence(tmpdir):
    # The special-case picker fails, because saving masked_array
    # not implemented, but it just delegates to the standard pickler.
    rnd = np.random.RandomState(0)
    a = rnd.random_sample(10)
    a = np.ma.masked_greater(a, 0.5)
    filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump(a, filename)
    b = numpy_pickle.load(filename, mmap_mode='r')
    assert isinstance(b, np.ma.masked_array)

@with_numpy
def test_compress_mmap_mode_warning(tmpdir):
    # Test the warning in case of compress + mmap_mode
    rnd = np.random.RandomState(0)
    a = rnd.random_sample(10)
    this_filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump(a, this_filename, compress=1)
    with warns(UserWarning) as warninfo:
        numpy_pickle.load(this_filename, mmap_mode='r+')
    debug_msg = "\n".join([str(w) for w in warninfo])
    warninfo = [w.message for w in warninfo]
    assert len(warninfo) == 1, debug_msg
    assert (
        str(warninfo[0]) ==
        'mmap_mode "r+" is not compatible with compressed '
        f'file {this_filename}. "r+" flag will be ignored.'
    )

@with_numpy
@parametrize('cache_size', [None, 0, 10])
def test_cache_size_warning(tmpdir, cache_size):
    # Check deprecation warning raised when cache size is not None
    filename = tmpdir.join('test.pkl').strpath
    rnd = np.random.RandomState(0)
    a = rnd.random_sample((10, 2))

    warnings.simplefilter("always")
    with warnings.catch_warnings(record=True) as warninfo:
        numpy_pickle.dump(a, filename, cache_size=cache_size)
    expected_nb_warnings = 1 if cache_size is not None else 0
    assert len(warninfo) == expected_nb_warnings
    for w in warninfo:
        assert w.category == DeprecationWarning
        assert (str(w.message) ==
                "Please do not set 'cache_size' in joblib.dump, this "
                "parameter has no effect and will be removed. You "
                "used 'cache_size={0}'".format(cache_size))

@with_numpy
@with_memory_profiler
@parametrize('compress', [True, False])
def test_memory_usage(tmpdir, compress):
    # Verify memory stays within expected bounds.
    filename = tmpdir.join('test.pkl').strpath
    small_array = np.ones((10, 10))
    big_array = np.ones(shape=100 * int(1e6), dtype=np.uint8)

    for obj in (small_array, big_array):
        size = obj.nbytes / 1e6
        obj_filename = filename + str(np.random.randint(0, 1000))
        mem_used = memory_used(numpy_pickle.dump,
                               obj, obj_filename, compress=compress)

        # The memory used to dump the object shouldn't exceed the buffer
        # size used to write array chunks (16MB).
        write_buf_size = _IO_BUFFER_SIZE + 16 * 1024 ** 2 / 1e6
        assert mem_used <= write_buf_size

        mem_used = memory_used(numpy_pickle.load, obj_filename)
        # memory used should be less than array size + buffer size used to
        # read the array chunk by chunk.
        read_buf_size = 32 + _IO_BUFFER_SIZE  # MiB
        assert mem_used < size + read_buf_size

@with_numpy
def test_compressed_pickle_dump_and_load(tmpdir):
    expected_list = [np.arange(5, dtype=np.dtype('<i8')),
                     np.arange(5, dtype=np.dtype('>i8')),
                     np.arange(5, dtype=np.dtype('<f8')),
                     np.arange(5, dtype=np.dtype('>f8')),
                     np.array([1, 'abc', {'a': 1, 'b': 2}], dtype='O'),
                     np.arange(256, dtype=np.uint8).tobytes(),
                     u"C'est l'\xe9t\xe9 !"]

    fname = tmpdir.join('temp.pkl.gz').strpath

    dumped_filenames = numpy_pickle.dump(expected_list, fname, compress=1)
    assert len(dumped_filenames) == 1
    result_list = numpy_pickle.load(fname)
    for result, expected in zip(result_list, expected_list):
        if isinstance(expected, np.ndarray):
            expected = _ensure_native_byte_order(expected)
            assert result.dtype == expected.dtype
            np.testing.assert_equal(result, expected)
        else:
            assert result == expected

def _check_pickle(filename, expected_list, mmap_mode=None):
    """Helper function to test joblib pickle content.

    Note: currently only pickles containing an iterable are supported
    by this function.
    """
    version_match = re.match(r'.+py(\d)(\d).+', filename)
    py_version_used_for_writing = int(version_match.group(1))

    py_version_to_default_pickle_protocol = {2: 2, 3: 3}
    pickle_reading_protocol = py_version_to_default_pickle_protocol.get(3, 4)
    pickle_writing_protocol = py_version_to_default_pickle_protocol.get(
        py_version_used_for_writing, 4)
    if pickle_reading_protocol >= pickle_writing_protocol:
        try:
            with warnings.catch_warnings(record=True) as warninfo:
                warnings.simplefilter('always')
                warnings.filterwarnings(
                    'ignore', module='numpy',
                    message='The compiler package is deprecated')
                result_list = numpy_pickle.load(filename, mmap_mode=mmap_mode)
            filename_base = os.path.basename(filename)
            expected_nb_deprecation_warnings = 1 if (
                "_0.9" in filename_base or "_0.8.4" in filename_base) else 0

            expected_nb_user_warnings = 3 if (
                re.search("_0.1.+.pkl$", filename_base) and
                mmap_mode is not None) else 0
            expected_nb_warnings = \
                expected_nb_deprecation_warnings + expected_nb_user_warnings
            assert len(warninfo) == expected_nb_warnings

            deprecation_warnings = [
                w for w in warninfo if issubclass(
                    w.category, DeprecationWarning)]
            user_warnings = [
                w for w in warninfo if issubclass(
                    w.category, UserWarning)]
            for w in deprecation_warnings:
                assert (str(w.message) ==
                        "The file '{0}' has been generated with a joblib "
                        "version less than 0.10. Please regenerate this "
                        "pickle file.".format(filename))

            for w in user_warnings:
                escaped_filename = re.escape(filename)
                assert re.search(
                    f"memmapped.+{escaped_filename}.+segmentation fault",
                    str(w.message))

            for result, expected in zip(result_list, expected_list):
                if isinstance(expected, np.ndarray):
                    expected = _ensure_native_byte_order(expected)
                    assert result.dtype == expected.dtype
                    np.testing.assert_equal(result, expected)
                else:
                    assert result == expected
        except Exception as exc:
            # When trying to read with python 3 a pickle generated
            # with python 2 we expect a user-friendly error
            if py_version_used_for_writing == 2:
                assert isinstance(exc, ValueError)
                message = ('You may be trying to read with '
                           'python 3 a joblib pickle generated with python 2.')
                assert message in str(exc)
            elif filename.endswith('.lz4') and with_lz4.args[0]:
                assert isinstance(exc, ValueError)
                assert LZ4_NOT_INSTALLED_ERROR in str(exc)
            else:
                raise
    else:
        # Pickle protocol used for writing is too high. We expect a
        # "unsupported pickle protocol" error message
        try:
            numpy_pickle.load(filename)
            raise AssertionError('Numpy pickle loading should '
                                 'have raised a ValueError exception')
        except ValueError as e:
            message = 'unsupported pickle protocol: {0}'.format(
                pickle_writing_protocol)
            assert message in str(e.args)

@with_numpy
def test_joblib_pickle_across_python_versions():
    # We need to be specific about dtypes in particular endianness
    # because the pickles can be generated on one architecture and
    # the tests run on another one. See
    # https://github.com/joblib/joblib/issues/279.
    expected_list = [np.arange(5, dtype=np.dtype('<i8')),
                     np.arange(5, dtype=np.dtype('<f8')),
                     np.array([1, 'abc', {'a': 1, 'b': 2}], dtype='O'),
                     np.arange(256, dtype=np.uint8).tobytes(),
                     # np.matrix is a subclass of np.ndarray, here we want
                     # to verify this type of object is correctly unpickled
                     # among versions.
                     np.matrix([0, 1, 2], dtype=np.dtype('<i8')),
                     u"C'est l'\xe9t\xe9 !"]

    # Testing all the compressed and non compressed
    # pickles in joblib/test/data. These pickles were generated by
    # the joblib/test/data/create_numpy_pickle.py script for the
    # relevant python, joblib and numpy versions.
    test_data_dir = os.path.dirname(os.path.abspath(data.__file__))

    pickle_extensions = ('.pkl', '.gz', '.gzip', '.bz2', 'lz4')
    if lzma is not None:
        pickle_extensions += ('.xz', '.lzma')
    pickle_filenames = [os.path.join(test_data_dir, fn)
                        for fn in os.listdir(test_data_dir)
                        if any(fn.endswith(ext) for ext in pickle_extensions)]

    for fname in pickle_filenames:
        _check_pickle(fname, expected_list)

@with_numpy
def test_joblib_pickle_across_python_versions_with_mmap():
    expected_list = [np.arange(5, dtype=np.dtype('<i8')),
                     np.arange(5, dtype=np.dtype('<f8')),
                     np.array([1, 'abc', {'a': 1, 'b': 2}], dtype='O'),
                     np.arange(256, dtype=np.uint8).tobytes(),
                     # np.matrix is a subclass of np.ndarray, here we want
                     # to verify this type of object is correctly unpickled
                     # among versions.
                     np.matrix([0, 1, 2], dtype=np.dtype('<i8')),
                     u"C'est l'\xe9t\xe9 !"]

    test_data_dir = os.path.dirname(os.path.abspath(data.__file__))

    pickle_filenames = [
        os.path.join(test_data_dir, fn)
        for fn in os.listdir(test_data_dir) if fn.endswith('.pkl')]
    for fname in pickle_filenames:
        _check_pickle(fname, expected_list, mmap_mode='r')

@with_numpy
def test_numpy_array_byte_order_mismatch_detection():
    # List of numpy arrays with big endian byteorder.
    be_arrays = [np.array([(1, 2.0), (3, 4.0)],
                          dtype=[('', '>i8'), ('', '>f8')]),
                 np.arange(3, dtype=np.dtype('>i8')),
                 np.arange(3, dtype=np.dtype('>f8'))]

    # Verify the byteorder mismatch is correctly detected.
    for array in be_arrays:
        if sys.byteorder == 'big':
            assert not _is_numpy_array_byte_order_mismatch(array)
        else:
            assert _is_numpy_array_byte_order_mismatch(array)
        converted = _ensure_native_byte_order(array)
        if converted.dtype.fields:
            for f in converted.dtype.fields.values():
                f[0].byteorder == '='
        else:
            assert converted.dtype.byteorder == "="

    # List of numpy arrays with little endian byteorder.
    le_arrays = [np.array([(1, 2.0), (3, 4.0)],
                          dtype=[('', '<i8'), ('', '<f8')]),
                 np.arange(3, dtype=np.dtype('<i8')),
                 np.arange(3, dtype=np.dtype('<f8'))]

    # Verify the byteorder mismatch is correctly detected.
    for array in le_arrays:
        if sys.byteorder == 'little':
            assert not _is_numpy_array_byte_order_mismatch(array)
        else:
            assert _is_numpy_array_byte_order_mismatch(array)
        converted = _ensure_native_byte_order(array)
        if converted.dtype.fields:
            for f in converted.dtype.fields.values():
                f[0].byteorder == '='
        else:
            assert converted.dtype.byteorder == "="

@parametrize('compress_tuple', [('zlib', 3), ('gzip', 3)])
def test_compress_tuple_argument(tmpdir, compress_tuple):
    # Verify the tuple is correctly taken into account.
    filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump("dummy", filename,
                      compress=compress_tuple)
    # Verify the file contains the right magic number
    with open(filename, 'rb') as f:
        assert _detect_compressor(f) == compress_tuple[0]

@parametrize('compress_tuple,message',
             [(('zlib', 3, 'extra'),        # wrong compress tuple
               'Compress argument tuple should contain exactly 2 elements'),
              (('wrong', 3),                # wrong compress method
               'Non valid compression method given: "{}"'.format('wrong')),
              (('zlib', 'wrong'),           # wrong compress level
               'Non valid compress level given: "{}"'.format('wrong'))])
def test_compress_tuple_argument_exception(tmpdir, compress_tuple, message):
    filename = tmpdir.join('test.pkl').strpath
    # Verify setting a wrong compress tuple raises a ValueError.
    with raises(ValueError) as excinfo:
        numpy_pickle.dump('dummy', filename, compress=compress_tuple)
    excinfo.match(message)

@parametrize('compress_string', ['zlib', 'gzip'])
def test_compress_string_argument(tmpdir, compress_string):
    # Verify the string is correctly taken into account.
    filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump("dummy", filename,
                      compress=compress_string)
    # Verify the file contains the right magic number
    with open(filename, 'rb') as f:
        assert _detect_compressor(f) == compress_string

@with_numpy
@parametrize('compress', [1, 3, 6])
@parametrize('cmethod', _COMPRESSORS)
def test_joblib_compression_formats(tmpdir, compress, cmethod):
    filename = tmpdir.join('test.pkl').strpath
    objects = (np.ones(shape=(100, 100), dtype='f8'),
               range(10),
               {'a': 1, 2: 'b'}, [], (), {}, 0, 1.0)

    if cmethod in ("lzma", "xz") and lzma is None:
        pytest.skip("lzma is support not available")

    elif cmethod == 'lz4' and with_lz4.args[0]:
        # Skip the test if lz4 is not installed. We here use the with_lz4
        # skipif fixture whose argument is True when lz4 is not installed
        pytest.skip("lz4 is not installed.")

    dump_filename = filename + "." + cmethod
    for obj in objects:
        numpy_pickle.dump(obj, dump_filename, compress=(cmethod, compress))
        # Verify the file contains the right magic number
        with open(dump_filename, 'rb') as f:
            assert _detect_compressor(f) == cmethod
        # Verify the reloaded object is correct
        obj_reloaded = numpy_pickle.load(dump_filename)
        assert isinstance(obj_reloaded, type(obj))
        if isinstance(obj, np.ndarray):
            np.testing.assert_array_equal(obj_reloaded, obj)
        else:
            assert obj_reloaded == obj

def _gzip_file_decompress(source_filename, target_filename):
    """Decompress a gzip file."""
    with closing(gzip.GzipFile(source_filename, "rb")) as fo:
        buf = fo.read()

    with open(target_filename, "wb") as fo:
        fo.write(buf)

def _zlib_file_decompress(source_filename, target_filename):
    """Decompress a zlib file."""
    with open(source_filename, 'rb') as fo:
        buf = zlib.decompress(fo.read())

    with open(target_filename, 'wb') as fo:
        fo.write(buf)

@parametrize('extension,decompress',
             [('.z', _zlib_file_decompress),
              ('.gz', _gzip_file_decompress)])
def test_load_externally_decompressed_files(tmpdir, extension, decompress):
    # Test that BinaryZlibFile generates valid gzip and zlib compressed files.
    obj = "a string to persist"
    filename_raw = tmpdir.join('test.pkl').strpath

    filename_compressed = filename_raw + extension
    # Use automatic extension detection to compress with the right method.
    numpy_pickle.dump(obj, filename_compressed)

    # Decompress with the corresponding method
    decompress(filename_compressed, filename_raw)

    # Test that the uncompressed pickle can be loaded and
    # that the result is correct.
    obj_reloaded = numpy_pickle.load(filename_raw)
    assert obj == obj_reloaded

@parametrize('extension,cmethod',
             # valid compressor extensions
             [('.z', 'zlib'),
              ('.gz', 'gzip'),
              ('.bz2', 'bz2'),
              ('.lzma', 'lzma'),
              ('.xz', 'xz'),
              # invalid compressor extensions
              ('.pkl', 'not-compressed'),
              ('', 'not-compressed')])
def test_compression_using_file_extension(tmpdir, extension, cmethod):
    if cmethod in ("lzma", "xz") and lzma is None:
        pytest.skip("lzma is missing")
    # test that compression method corresponds to the given filename extension.
    filename = tmpdir.join('test.pkl').strpath
    obj = "object to dump"

    dump_fname = filename + extension
    numpy_pickle.dump(obj, dump_fname)
    # Verify the file contains the right magic number
    with open(dump_fname, 'rb') as f:
        assert _detect_compressor(f) == cmethod
    # Verify the reloaded object is correct
    obj_reloaded = numpy_pickle.load(dump_fname)
    assert isinstance(obj_reloaded, type(obj))
    assert obj_reloaded == obj

@with_numpy
def test_file_handle_persistence(tmpdir):
    objs = [np.random.random((10, 10)), "some data"]
    fobjs = [bz2.BZ2File, gzip.GzipFile]
    if lzma is not None:
        fobjs += [lzma.LZMAFile]
    filename = tmpdir.join('test.pkl').strpath

    for obj in objs:
        for fobj in fobjs:
            with fobj(filename, 'wb') as f:
                numpy_pickle.dump(obj, f)

            # using the same decompressor prevents from internally
            # decompress again.
            with fobj(filename, 'rb') as f:
                obj_reloaded = numpy_pickle.load(f)

            # when needed, the correct decompressor should be used when
            # passing a raw file handle.
            with open(filename, 'rb') as f:
                obj_reloaded_2 = numpy_pickle.load(f)

            if isinstance(obj, np.ndarray):
                np.testing.assert_array_equal(obj_reloaded, obj)
                np.testing.assert_array_equal(obj_reloaded_2, obj)
            else:
                assert obj_reloaded == obj
                assert obj_reloaded_2 == obj

@with_numpy
def test_in_memory_persistence():
    objs = [np.random.random((10, 10)), "some data"]
    for obj in objs:
        f = io.BytesIO()
        numpy_pickle.dump(obj, f)
        obj_reloaded = numpy_pickle.load(f)
        if isinstance(obj, np.ndarray):
            np.testing.assert_array_equal(obj_reloaded, obj)
        else:
            assert obj_reloaded == obj

@with_numpy
def test_file_handle_persistence_mmap(tmpdir):
    obj = np.random.random((10, 10))
    filename = tmpdir.join('test.pkl').strpath

    with open(filename, 'wb') as f:
        numpy_pickle.dump(obj, f)

    with open(filename, 'rb') as f:
        obj_reloaded = numpy_pickle.load(f, mmap_mode='r+')

    np.testing.assert_array_equal(obj_reloaded, obj)

@with_numpy
def test_file_handle_persistence_compressed_mmap(tmpdir):
    obj = np.random.random((10, 10))
    filename = tmpdir.join('test.pkl').strpath

    with open(filename, 'wb') as f:
        numpy_pickle.dump(obj, f, compress=('gzip', 3))

    with closing(gzip.GzipFile(filename, 'rb')) as f:
        with warns(UserWarning) as warninfo:
            numpy_pickle.load(f, mmap_mode='r+')
        assert len(warninfo) == 1
        assert (str(warninfo[0].message) ==
                '"%(fileobj)r" is not a raw file, mmap_mode "%(mmap_mode)s" '
                'flag will be ignored.' % {'fileobj': f, 'mmap_mode': 'r+'})

@with_numpy
def test_file_handle_persistence_in_memory_mmap():
    obj = np.random.random((10, 10))
    buf = io.BytesIO()

    numpy_pickle.dump(obj, buf)

    with warns(UserWarning) as warninfo:
        numpy_pickle.load(buf, mmap_mode='r+')
    assert len(warninfo) == 1
    assert (str(warninfo[0].message) ==
            'In memory persistence is not compatible with mmap_mode '
            '"%(mmap_mode)s" flag passed. mmap_mode option will be '
            'ignored.' % {'mmap_mode': 'r+'})

@parametrize('data', [b'a little data as bytes.',
                      # More bytes
                      10000 * "{}".format(
                          random.randint(0, 1000) * 1000).encode('latin-1')],
             ids=["a little data as bytes.", "a large data as bytes."])
@parametrize('compress_level', [1, 3, 9])
def test_binary_zlibfile(tmpdir, data, compress_level):
    filename = tmpdir.join('test.pkl').strpath
    # Regular cases
    with open(filename, 'wb') as f:
        with BinaryZlibFile(f, 'wb',
                            compresslevel=compress_level) as fz:
            assert fz.writable()
            fz.write(data)
            assert fz.fileno() == f.fileno()
            with raises(io.UnsupportedOperation):
                fz._check_can_read()

            with raises(io.UnsupportedOperation):
                fz._check_can_seek()
        assert fz.closed
        with raises(ValueError):
            fz._check_not_closed()

    with open(filename, 'rb') as f:
        with BinaryZlibFile(f) as fz:
            assert fz.readable()
            assert fz.seekable()
            assert fz.fileno() == f.fileno()
            assert fz.read() == data
            with raises(io.UnsupportedOperation):
                fz._check_can_write()
            assert fz.seekable()
            fz.seek(0)
            assert fz.tell() == 0
        assert fz.closed

    # Test with a filename as input
    with BinaryZlibFile(filename, 'wb',
                        compresslevel=compress_level) as fz:
        assert fz.writable()
        fz.write(data)

    with BinaryZlibFile(filename, 'rb') as fz:
        assert fz.read() == data
        assert fz.seekable()

    # Test without context manager
    fz = BinaryZlibFile(filename, 'wb', compresslevel=compress_level)
    assert fz.writable()
    fz.write(data)
    fz.close()

    fz = BinaryZlibFile(filename, 'rb')
    assert fz.read() == data
    fz.close()

@parametrize('bad_value', [-1, 10, 15, 'a', (), {}])
def test_binary_zlibfile_bad_compression_levels(tmpdir, bad_value):
    filename = tmpdir.join('test.pkl').strpath
    with raises(ValueError) as excinfo:
        BinaryZlibFile(filename, 'wb', compresslevel=bad_value)
    pattern = re.escape("'compresslevel' must be an integer between 1 and 9. "
                        "You provided 'compresslevel={}'".format(bad_value))
    excinfo.match(pattern)

@parametrize('bad_mode', ['a', 'x', 'r', 'w', 1, 2])
def test_binary_zlibfile_invalid_modes(tmpdir, bad_mode):
    filename = tmpdir.join('test.pkl').strpath
    with raises(ValueError) as excinfo:
        BinaryZlibFile(filename, bad_mode)
    excinfo.match("Invalid mode")

@parametrize('bad_file', [1, (), {}])
def test_binary_zlibfile_invalid_filename_type(bad_file):
    with raises(TypeError) as excinfo:
        BinaryZlibFile(bad_file, 'rb')
    excinfo.match("filename must be a str or bytes object, or a file")

###############################################################################
# Test dumping array subclasses
if np is not None:

    class SubArray(np.ndarray):

        def __reduce__(self):
            return _load_sub_array, (np.asarray(self), )

    def _load_sub_array(arr):
        d = SubArray(arr.shape)
        d[:] = arr
        return d

    class ComplexTestObject:
        """A complex object containing numpy arrays as attributes."""

        def __init__(self):
            self.array_float = np.arange(100, dtype='float64')
            self.array_int = np.ones(100, dtype='int32')
            self.array_obj = np.array(['a', 10, 20.0], dtype='object')

@with_numpy
def test_numpy_subclass(tmpdir):
    filename = tmpdir.join('test.pkl').strpath
    a = SubArray((10,))
    numpy_pickle.dump(a, filename)
    c = numpy_pickle.load(filename)
    assert isinstance(c, SubArray)
    np.testing.assert_array_equal(c, a)

def test_pathlib(tmpdir):
    filename = tmpdir.join('test.pkl').strpath
    value = 123
    numpy_pickle.dump(value, Path(filename))
    assert numpy_pickle.load(filename) == value
    numpy_pickle.dump(value, filename)
    assert numpy_pickle.load(Path(filename)) == value

@with_numpy
def test_non_contiguous_array_pickling(tmpdir):
    filename = tmpdir.join('test.pkl').strpath

    for array in [  # Array that triggers a contiguousness issue with nditer,
                    # see https://github.com/joblib/joblib/pull/352 and see
                    # https://github.com/joblib/joblib/pull/353
                    np.asfortranarray([[1, 2], [3, 4]])[1:],
                    # Non contiguous array with works fine with nditer
                    np.ones((10, 50, 20), order='F')[:, :1, :]]:
        assert not array.flags.c_contiguous
        assert not array.flags.f_contiguous
        numpy_pickle.dump(array, filename)
        array_reloaded = numpy_pickle.load(filename)
        np.testing.assert_array_equal(array_reloaded, array)

@with_numpy
def test_pickle_highest_protocol(tmpdir):
    # ensure persistence of a numpy array is valid even when using
    # the pickle HIGHEST_PROTOCOL.
    # see https://github.com/joblib/joblib/issues/362

    filename = tmpdir.join('test.pkl').strpath
    test_array = np.zeros(10)

    numpy_pickle.dump(test_array, filename, protocol=pickle.HIGHEST_PROTOCOL)
    array_reloaded = numpy_pickle.load(filename)

    np.testing.assert_array_equal(array_reloaded, test_array)

@with_numpy
def test_pickle_in_socket():
    # test that joblib can pickle in sockets
    test_array = np.arange(10)
    _ADDR = ("localhost", 12345)
    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    listener.bind(_ADDR)
    listener.listen(1)

    with socket.create_connection(_ADDR) as client:
        server, client_addr = listener.accept()

        with server.makefile("wb") as sf:
            numpy_pickle.dump(test_array, sf)

        with client.makefile("rb") as cf:
            array_reloaded = numpy_pickle.load(cf)

        np.testing.assert_array_equal(array_reloaded, test_array)

        # Check that a byte-aligned numpy array written in a file can be send
        # over a socket and then read on the other side
        bytes_to_send = io.BytesIO()
        numpy_pickle.dump(test_array, bytes_to_send)
        server.send(bytes_to_send.getvalue())

        with client.makefile("rb") as cf:
            array_reloaded = numpy_pickle.load(cf)

        np.testing.assert_array_equal(array_reloaded, test_array)

@with_numpy
def test_load_memmap_with_big_offset(tmpdir):
    # Test that numpy memmap offset is set correctly if greater than
    # mmap.ALLOCATIONGRANULARITY, see
    # https://github.com/joblib/joblib/issues/451 and
    # https://github.com/numpy/numpy/pull/8443 for more details.
    fname = tmpdir.join('test.mmap').strpath
    size = mmap.ALLOCATIONGRANULARITY
    obj = [np.zeros(size, dtype='uint8'), np.ones(size, dtype='uint8')]
    numpy_pickle.dump(obj, fname)
    memmaps = numpy_pickle.load(fname, mmap_mode='r')
    assert isinstance(memmaps[1], np.memmap)
    assert memmaps[1].offset > size
    np.testing.assert_array_equal(obj, memmaps)

def test_register_compressor(tmpdir):
    # Check that registering compressor file works.
    compressor_name = 'test-name'
    compressor_prefix = 'test-prefix'

    class BinaryCompressorTestFile(io.BufferedIOBase):
        pass

    class BinaryCompressorTestWrapper(CompressorWrapper):

        def __init__(self):
            CompressorWrapper.__init__(self, obj=BinaryCompressorTestFile,
                                       prefix=compressor_prefix)

    register_compressor(compressor_name, BinaryCompressorTestWrapper())

    assert (_COMPRESSORS[compressor_name].fileobj_factory ==
            BinaryCompressorTestFile)
    assert _COMPRESSORS[compressor_name].prefix == compressor_prefix

    # Remove this dummy compressor file from extra compressors because other
    # tests might fail because of this
    _COMPRESSORS.pop(compressor_name)

@parametrize('invalid_name', [1, (), {}])
def test_register_compressor_invalid_name(invalid_name):
    # Test that registering an invalid compressor name is not allowed.
    with raises(ValueError) as excinfo:
        register_compressor(invalid_name, None)
    excinfo.match("Compressor name should be a string")

def test_register_compressor_invalid_fileobj():
    # Test that registering an invalid file object is not allowed.

    class InvalidFileObject():
        pass

    class InvalidFileObjectWrapper(CompressorWrapper):
        def __init__(self):
            CompressorWrapper.__init__(self, obj=InvalidFileObject,
                                       prefix=b'prefix')

    with raises(ValueError) as excinfo:
        register_compressor('invalid', InvalidFileObjectWrapper())

    excinfo.match("Compressor 'fileobj_factory' attribute should implement "
                  "the file object interface")

class AnotherZlibCompressorWrapper(CompressorWrapper):

    def __init__(self):
        CompressorWrapper.__init__(self, obj=BinaryZlibFile, prefix=b'prefix')

class StandardLibGzipCompressorWrapper(CompressorWrapper):

    def __init__(self):
        CompressorWrapper.__init__(self, obj=gzip.GzipFile, prefix=b'prefix')

def test_register_compressor_already_registered():
    # Test registration of existing compressor files.
    compressor_name = 'test-name'

    # register a test compressor
    register_compressor(compressor_name, AnotherZlibCompressorWrapper())

    with raises(ValueError) as excinfo:
        register_compressor(compressor_name,
                            StandardLibGzipCompressorWrapper())
    excinfo.match("Compressor '{}' already registered."
                  .format(compressor_name))

    register_compressor(compressor_name, StandardLibGzipCompressorWrapper(),
                        force=True)

    assert compressor_name in _COMPRESSORS
    assert _COMPRESSORS[compressor_name].fileobj_factory == gzip.GzipFile

    # Remove this dummy compressor file from extra compressors because other
    # tests might fail because of this
    _COMPRESSORS.pop(compressor_name)

@with_lz4
def test_lz4_compression(tmpdir):
    # Check that lz4 can be used when dependency is available.
    import lz4.frame
    compressor = 'lz4'
    assert compressor in _COMPRESSORS
    assert _COMPRESSORS[compressor].fileobj_factory == lz4.frame.LZ4FrameFile

    fname = tmpdir.join('test.pkl').strpath
    data = 'test data'
    numpy_pickle.dump(data, fname, compress=compressor)

    with open(fname, 'rb') as f:
        assert f.read(len(_LZ4_PREFIX)) == _LZ4_PREFIX
    assert numpy_pickle.load(fname) == data

    # Test that LZ4 is applied based on file extension
    numpy_pickle.dump(data, fname + '.lz4')
    with open(fname, 'rb') as f:
        assert f.read(len(_LZ4_PREFIX)) == _LZ4_PREFIX
    assert numpy_pickle.load(fname) == data

@without_lz4
def test_lz4_compression_without_lz4(tmpdir):
    # Check that lz4 cannot be used when dependency is not available.
    fname = tmpdir.join('test.nolz4').strpath
    data = 'test data'
    msg = LZ4_NOT_INSTALLED_ERROR
    with raises(ValueError) as excinfo:
        numpy_pickle.dump(data, fname, compress='lz4')
    excinfo.match(msg)

    with raises(ValueError) as excinfo:
        numpy_pickle.dump(data, fname + '.lz4')
    excinfo.match(msg)

protocols = [pickle.DEFAULT_PROTOCOL]
if pickle.HIGHEST_PROTOCOL != pickle.DEFAULT_PROTOCOL:
    protocols.append(pickle.HIGHEST_PROTOCOL)

@with_numpy
@parametrize('protocol', protocols)
def test_memmap_alignment_padding(tmpdir, protocol):
    # Test that memmaped arrays returned by numpy.load are correctly aligned
    fname = tmpdir.join('test.mmap').strpath

    a = np.random.randn(2)
    numpy_pickle.dump(a, fname, protocol=protocol)
    memmap = numpy_pickle.load(fname, mmap_mode='r')
    assert isinstance(memmap, np.memmap)
    np.testing.assert_array_equal(a, memmap)
    assert (
        memmap.ctypes.data % numpy_pickle.NUMPY_ARRAY_ALIGNMENT_BYTES == 0)
    assert memmap.flags.aligned

    array_list = [
        np.random.randn(2), np.random.randn(2),
        np.random.randn(2), np.random.randn(2)
    ]

    # On Windows OSError 22 if reusing the same path for memmap ...
    fname = tmpdir.join('test1.mmap').strpath
    numpy_pickle.dump(array_list, fname, protocol=protocol)
    l_reloaded = numpy_pickle.load(fname, mmap_mode='r')

    for idx, memmap in enumerate(l_reloaded):
        assert isinstance(memmap, np.memmap)
        np.testing.assert_array_equal(array_list[idx], memmap)
        assert (
            memmap.ctypes.data % numpy_pickle.NUMPY_ARRAY_ALIGNMENT_BYTES == 0)
        assert memmap.flags.aligned

    array_dict = {
        'a0': np.arange(2, dtype=np.uint8),
        'a1': np.arange(3, dtype=np.uint8),
        'a2': np.arange(5, dtype=np.uint8),
        'a3': np.arange(7, dtype=np.uint8),
        'a4': np.arange(11, dtype=np.uint8),
        'a5': np.arange(13, dtype=np.uint8),
        'a6': np.arange(17, dtype=np.uint8),
        'a7': np.arange(19, dtype=np.uint8),
        'a8': np.arange(23, dtype=np.uint8),
    }

    # On Windows OSError 22 if reusing the same path for memmap ...
    fname = tmpdir.join('test2.mmap').strpath
    numpy_pickle.dump(array_dict, fname, protocol=protocol)
    d_reloaded = numpy_pickle.load(fname, mmap_mode='r')

    for key, memmap in d_reloaded.items():
        assert isinstance(memmap, np.memmap)
        np.testing.assert_array_equal(array_dict[key], memmap)
        assert (
            memmap.ctypes.data % numpy_pickle.NUMPY_ARRAY_ALIGNMENT_BYTES == 0)
        assert memmap.flags.aligned

"""
===============================
Improving I/O using compressors
===============================

This example compares the compressors available in Joblib. In the example,
Zlib, LZMA and LZ4 compression only are used but Joblib also supports BZ2 and
GZip compression methods.
For each compared compression method, this example dumps and reloads a
dataset fetched from an online machine-learning database. This gives 3
information: the size on disk of the compressed data, the time spent to dump
and the time spent to reload the data from disk.
"""

import os
import os.path
import time

###############################################################################
# Get some data from real-world use cases
# ---------------------------------------
#
# First fetch the benchmark dataset from an online machine-learning database
# and load it in a pandas dataframe.

import pandas as pd

url = "https://github.com/joblib/dataset/raw/main/kddcup.data.gz"
names = ("duration, protocol_type, service, flag, src_bytes, "
         "dst_bytes, land, wrong_fragment, urgent, hot, "
         "num_failed_logins, logged_in, num_compromised, "
         "root_shell, su_attempted, num_root, "
         "num_file_creations, ").split(', ')

data = pd.read_csv(url, names=names, nrows=1e6)

###############################################################################
# Dump and load the dataset without compression
# ---------------------------------------------
#
# This gives reference values for later comparison.

from joblib import dump, load

pickle_file = './pickle_data.joblib'

###############################################################################
# Start by measuring the time spent for dumping the raw data:
start = time.time()
with open(pickle_file, 'wb') as f:
    dump(data, f)
raw_dump_duration = time.time() - start
print("Raw dump duration: %0.3fs" % raw_dump_duration)

###############################################################################
# Then measure the size of the raw dumped data on disk:
raw_file_size = os.stat(pickle_file).st_size / 1e6
print("Raw dump file size: %0.3fMB" % raw_file_size)

###############################################################################
# Finally measure the time spent for loading the raw data:
start = time.time()
with open(pickle_file, 'rb') as f:
    load(f, allow_pickle=True)
raw_load_duration = time.time() - start
print("Raw load duration: %0.3fs" % raw_load_duration)

###############################################################################
# Dump and load the dataset using the Zlib compression method
# -----------------------------------------------------------
#
# The compression level is using the default value, 3, which is, in general, a
# good compromise between compression and speed.

###############################################################################
# Start by measuring the time spent for dumping of the zlib data:

start = time.time()
with open(pickle_file, 'wb') as f:
    dump(data, f, compress='zlib')
zlib_dump_duration = time.time() - start
print("Zlib dump duration: %0.3fs" % zlib_dump_duration)

###############################################################################
# Then measure the size of the zlib dump data on disk:

zlib_file_size = os.stat(pickle_file).st_size / 1e6
print("Zlib file size: %0.3fMB" % zlib_file_size)

###############################################################################
# Finally measure the time spent for loading the compressed dataset:

start = time.time()
with open(pickle_file, 'rb') as f:
    load(f, allow_pickle=True)
zlib_load_duration = time.time() - start
print("Zlib load duration: %0.3fs" % zlib_load_duration)

###############################################################################
# .. note:: The compression format is detected automatically by Joblib.
#    The compression format is identified by the standard magic number present
#    at the beginning of the file. Joblib uses this information to determine
#    the compression method used.
#    This is the case for all compression methods supported by Joblib.

###############################################################################
# Dump and load the dataset using the LZMA compression method
# -----------------------------------------------------------
#
# LZMA compression method has a very good compression rate but at the cost
# of being very slow.
# In this example, a light compression level, e.g. 3, is used to speed up a
# bit the dump/load cycle.

###############################################################################
# Start by measuring the time spent for dumping the lzma data:

start = time.time()
with open(pickle_file, 'wb') as f:
    dump(data, f, compress=('lzma', 3))
lzma_dump_duration = time.time() - start
print("LZMA dump duration: %0.3fs" % lzma_dump_duration)

###############################################################################
# Then measure the size of the lzma dump data on disk:

lzma_file_size = os.stat(pickle_file).st_size / 1e6
print("LZMA file size: %0.3fMB" % lzma_file_size)

###############################################################################
# Finally measure the time spent for loading the lzma data:

start = time.time()
with open(pickle_file, 'rb') as f:
    load(f, allow_pickle=True)
lzma_load_duration = time.time() - start
print("LZMA load duration: %0.3fs" % lzma_load_duration)

###############################################################################
# Dump and load the dataset using the LZ4 compression method
# ----------------------------------------------------------
#
# LZ4 compression method is known to be one of the fastest available
# compression method but with a compression rate a bit lower than Zlib. In
# most of the cases, this method is a good choice.

###############################################################################
# .. note:: In order to use LZ4 compression with Joblib, the
#    `lz4 <https://pypi.python.org/pypi/lz4>`_ package must be installed
#    on the system.

###############################################################################
# Start by measuring the time spent for dumping the lz4 data:

start = time.time()
with open(pickle_file, 'wb') as f:
    dump(data, f, compress='lz4')
lz4_dump_duration = time.time() - start
print("LZ4 dump duration: %0.3fs" % lz4_dump_duration)

###############################################################################
# Then measure the size of the lz4 dump data on disk:

lz4_file_size = os.stat(pickle_file).st_size / 1e6
print("LZ4 file size: %0.3fMB" % lz4_file_size)

###############################################################################
# Finally measure the time spent for loading the lz4 data:

start = time.time()
with open(pickle_file, 'rb') as f:
    load(f, allow_pickle=True)
lz4_load_duration = time.time() - start
print("LZ4 load duration: %0.3fs" % lz4_load_duration)

###############################################################################
# Comparing the results
# ---------------------

import numpy as np
import matplotlib.pyplot as plt

N = 4
load_durations = (raw_load_duration, lz4_load_duration, zlib_load_duration,
                  lzma_load_duration)
dump_durations = (raw_dump_duration, lz4_dump_duration, zlib_dump_duration,
                  lzma_dump_duration)
file_sizes = (raw_file_size, lz4_file_size, zlib_file_size, lzma_file_size)
ind = np.arange(N)
width = 0.5

plt.figure(1, figsize=(5, 4))
p1 = plt.bar(ind, dump_durations, width)
p2 = plt.bar(ind, load_durations, width, bottom=dump_durations)
plt.ylabel('Time in seconds')
plt.title('Dump and load durations')
plt.xticks(ind, ('Raw', 'LZ4', 'Zlib', 'LZMA'))
plt.yticks(np.arange(0, lzma_load_duration + lzma_dump_duration))
plt.legend((p1[0], p2[0]), ('Dump duration', 'Load duration'))

###############################################################################
# Compared with other compressors, LZ4 is clearly the fastest, especially for
# dumping compressed data on disk. In this particular case, it can even be
# faster than the raw dump.
# Also note that dump and load durations depend on the I/O speed of the
# underlying storage: for example, with SSD hard drives the LZ4 compression
# will be slightly slower than raw dump/load, whereas with spinning hard disk
# drives (HDD) or remote storage (NFS), LZ4 is faster in general.
#
# LZMA and Zlib, even if always slower for dumping data, are quite fast when
# re-loading compressed data from disk.

plt.figure(2, figsize=(5, 4))
plt.bar(ind, file_sizes, width, log=True)
plt.ylabel('File size in MB')
plt.xticks(ind, ('Raw', 'LZ4', 'Zlib', 'LZMA'))

###############################################################################
# Compressed data obviously takes a lot less space on disk than raw data. LZMA
# is the best compression method in terms of compression rate. Zlib also has a
# better compression rate than LZ4.

plt.show()

###############################################################################
# Clear the pickle file
# ---------------------

import os
os.remove(pickle_file)

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

# Author: Gael Varoquaux <gael dot varoquaux at normalesup dot org>
# Copyright (c) 2009 Gael Varoquaux
# License: BSD Style, 3 clauses.

import pickle
import os
import warnings
import io
from pathlib import Path

from .compressor import lz4, LZ4_NOT_INSTALLED_ERROR
from .compressor import _COMPRESSORS, register_compressor, BinaryZlibFile
from .compressor import (ZlibCompressorWrapper, GzipCompressorWrapper,
                         BZ2CompressorWrapper, LZMACompressorWrapper,
                         XZCompressorWrapper, LZ4CompressorWrapper)
from .numpy_pickle_utils import Unpickler, Pickler
from .numpy_pickle_utils import _read_fileobject, _write_fileobject
from .numpy_pickle_utils import _read_bytes, BUFFER_SIZE
from .numpy_pickle_utils import _ensure_native_byte_order
from .numpy_pickle_compat import load_compatibility
from .numpy_pickle_compat import NDArrayWrapper
# For compatibility with old versions of joblib, we need ZNDArrayWrapper
# to be visible in the current namespace.
# Explicitly skipping next line from flake8 as it triggers an F401 warning
# which we don't care.
from .numpy_pickle_compat import ZNDArrayWrapper  # noqa
from .backports import make_memmap

# Register supported compressors
register_compressor('zlib', ZlibCompressorWrapper())
register_compressor('gzip', GzipCompressorWrapper())
register_compressor('bz2', BZ2CompressorWrapper())
register_compressor('lzma', LZMACompressorWrapper())
register_compressor('xz', XZCompressorWrapper())
register_compressor('lz4', LZ4CompressorWrapper())

###############################################################################
# Utility objects for persistence.

# For convenience, 16 bytes are used to be sure to cover all the possible
# dtypes' alignments. For reference, see:
# https://numpy.org/devdocs/dev/alignment.html
NUMPY_ARRAY_ALIGNMENT_BYTES = 16

class NumpyArrayWrapper(object):
    """An object to be persisted instead of numpy arrays.

    This object is used to hack into the pickle machinery and read numpy
    array data from our custom persistence format.
    More precisely, this object is used for:
    * carrying the information of the persisted array: subclass, shape, order,
    dtype. Those ndarray metadata are used to correctly reconstruct the array
    with low level numpy functions.
    * determining if memmap is allowed on the array.
    * reading the array bytes from a file.
    * reading the array using memorymap from a file.
    * writing the array bytes to a file.

    Attributes
    ----------
    subclass: numpy.ndarray subclass
        Determine the subclass of the wrapped array.
    shape: numpy.ndarray shape
        Determine the shape of the wrapped array.
    order: {'C', 'F'}
        Determine the order of wrapped array data. 'C' is for C order, 'F' is
        for fortran order.
    dtype: numpy.ndarray dtype
        Determine the data type of the wrapped array.
    allow_mmap: bool
        Determine if memory mapping is allowed on the wrapped array.
        Default: False.
    allow_pickle: bool
        Determine if pickle.load is allowed on the wrapped array.
        Default: False.
    """

    def __init__(self, subclass, shape, order, dtype,
                 allow_mmap=False, allow_pickle=False,
                 numpy_array_alignment_bytes=NUMPY_ARRAY_ALIGNMENT_BYTES):
        """Constructor. Store the useful information for later."""
        self.subclass = subclass
        self.shape = shape
        self.order = order
        self.dtype = dtype
        self.allow_mmap = allow_mmap
        self.allow_pickle = allow_pickle
        # We make numpy_array_alignment_bytes an instance attribute to allow us
        # to change our mind about the default alignment and still load the old
        # pickles (with the previous alignment) correctly
        self.numpy_array_alignment_bytes = numpy_array_alignment_bytes

    def safe_get_numpy_array_alignment_bytes(self):
        # NumpyArrayWrapper instances loaded from joblib <= 1.1 pickles don't
        # have an numpy_array_alignment_bytes attribute
        return getattr(self, 'numpy_array_alignment_bytes', None)

    def write_array(self, array, pickler):
        """Write array bytes to pickler file handle.

        This function is an adaptation of the numpy write_array function
        available in version 1.10.1 in numpy/lib/format.py.
        """
        # Set buffer size to 16 MiB to hide the Python loop overhead.
        buffersize = max(16 * 1024 ** 2 // array.itemsize, 1)
        if array.dtype.hasobject:
            # We contain Python objects so we cannot write out the data
            # directly. Instead, we will pickle it out with version 2 of the
            # pickle protocol.
            pickle.dump(array, pickler.file_handle, protocol=2)
        else:
            numpy_array_alignment_bytes = \
                self.safe_get_numpy_array_alignment_bytes()
            if numpy_array_alignment_bytes is not None:
                current_pos = pickler.file_handle.tell()
                pos_after_padding_byte = current_pos + 1
                padding_length = numpy_array_alignment_bytes - (
                    pos_after_padding_byte % numpy_array_alignment_bytes)
                # A single byte is written that contains the padding length in
                # bytes
                padding_length_byte = int.to_bytes(
                    padding_length, length=1, byteorder='little')
                pickler.file_handle.write(padding_length_byte)

                if padding_length != 0:
                    padding = b'\xff' * padding_length
                    pickler.file_handle.write(padding)

            for chunk in pickler.np.nditer(array,
                                           flags=['external_loop',
                                                  'buffered',
                                                  'zerosize_ok'],
                                           buffersize=buffersize,
                                           order=self.order):
                pickler.file_handle.write(chunk.tobytes('C'))

    def read_array(self, unpickler, allow_pickle=False):
        """Read array from unpickler file handle.

        This function is an adaptation of the numpy read_array function
        available in version 1.10.1 in numpy/lib/format.py.
        """
        if len(self.shape) == 0:
            count = 1
        else:
            # joblib issue #859: we cast the elements of self.shape to int64 to
            # prevent a potential overflow when computing their product.
            shape_int64 = [unpickler.np.int64(x) for x in self.shape]
            count = unpickler.np.multiply.reduce(shape_int64)
        # Now read the actual data.
        if self.dtype.hasobject:
            if not allow_pickle:
                raise ValueError("Object arrays cannot be loaded when "
                                 "allow_pickle=False")
            # The array contained Python objects. We need to unpickle the data.
            array = pickle.load(unpickler.file_handle)
        else:
            numpy_array_alignment_bytes = \
                self.safe_get_numpy_array_alignment_bytes()
            if numpy_array_alignment_bytes is not None:
                padding_byte = unpickler.file_handle.read(1)
                padding_length = int.from_bytes(
                    padding_byte, byteorder='little')
                if padding_length != 0:
                    unpickler.file_handle.read(padding_length)

            # This is not a real file. We have to read it the
            # memory-intensive way.
            # crc32 module fails on reads greater than 2 ** 32 bytes,
            # breaking large reads from gzip streams. Chunk reads to
            # BUFFER_SIZE bytes to avoid issue and reduce memory overhead
            # of the read. In non-chunked case count < max_read_count, so
            # only one read is performed.
            max_read_count = BUFFER_SIZE // min(BUFFER_SIZE,
                                                self.dtype.itemsize)

            array = unpickler.np.empty(count, dtype=self.dtype)
            for i in range(0, count, max_read_count):
                read_count = min(max_read_count, count - i)
                read_size = int(read_count * self.dtype.itemsize)
                data = _read_bytes(unpickler.file_handle,
                                   read_size, "array data")
                array[i:i + read_count] = \
                    unpickler.np.frombuffer(data, dtype=self.dtype,
                                            count=read_count)
                del data

            if self.order == 'F':
                array.shape = self.shape[::-1]
                array = array.transpose()
            else:
                array.shape = self.shape

        # Detect byte order mismatch and swap as needed.
        return _ensure_native_byte_order(array)

    def read_mmap(self, unpickler):
        """Read an array using numpy memmap."""
        current_pos = unpickler.file_handle.tell()
        offset = current_pos
        numpy_array_alignment_bytes = \
            self.safe_get_numpy_array_alignment_bytes()

        if numpy_array_alignment_bytes is not None:
            padding_byte = unpickler.file_handle.read(1)
            padding_length = int.from_bytes(padding_byte, byteorder='little')
            # + 1 is for the padding byte
            offset += padding_length + 1

        if unpickler.mmap_mode == 'w+':
            unpickler.mmap_mode = 'r+'

        marray = make_memmap(unpickler.filename,
                             dtype=self.dtype,
                             shape=self.shape,
                             order=self.order,
                             mode=unpickler.mmap_mode,
                             offset=offset)
        # update the offset so that it corresponds to the end of the read array
        unpickler.file_handle.seek(offset + marray.nbytes)

        if (numpy_array_alignment_bytes is None and
                current_pos % NUMPY_ARRAY_ALIGNMENT_BYTES != 0):
            message = (
                f'The memmapped array {marray} loaded from the file '
                f'{unpickler.file_handle.name} is not byte aligned. '
                'This may cause segmentation faults if this memmapped array '
                'is used in some libraries like BLAS or PyTorch. '
                'To get rid of this warning, regenerate your pickle file '
                'with joblib >= 1.2.0. '
                'See https://github.com/joblib/joblib/issues/563 '
                'for more details'
            )
            warnings.warn(message)

        return _ensure_native_byte_order(marray)

    def read(self, unpickler, allow_pickle=False):
        """Read the array corresponding to this wrapper.

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

        Parameters
        ----------
        unpickler: NumpyUnpickler

        Returns
        -------
        array: numpy.ndarray

        """
        # When requested, only use memmap mode if allowed.
        if unpickler.mmap_mode is not None and self.allow_mmap:
            array = self.read_mmap(unpickler)
        else:
            array = self.read_array(unpickler, allow_pickle)

        # Manage array subclass case
        if (hasattr(array, '__array_prepare__') and
            self.subclass not in (unpickler.np.ndarray,
                                  unpickler.np.memmap)):
            # We need to reconstruct another subclass
            new_array = unpickler.np.core.multiarray._reconstruct(
                self.subclass, (0,), 'b')
            return new_array.__array_prepare__(array)
        else:
            return array

###############################################################################
# Pickler classes

class NumpyPickler(Pickler):
    """A pickler to persist big data efficiently.

    The main features of this object are:
    * persistence of numpy arrays in a single file.
    * optional compression with a special care on avoiding memory copies.

    Attributes
    ----------
    fp: file
        File object handle used for serializing the input object.
    protocol: int, optional
        Pickle protocol used. Default is pickle.DEFAULT_PROTOCOL.
    """

    dispatch = Pickler.dispatch.copy()

    def __init__(self, fp, protocol=None):
        self.file_handle = fp
        self.buffered = isinstance(self.file_handle, BinaryZlibFile)

        # By default we want a pickle protocol that only changes with
        # the major python version and not the minor one
        if protocol is None:
            protocol = pickle.DEFAULT_PROTOCOL

        Pickler.__init__(self, self.file_handle, protocol=protocol)
        # delayed import of numpy, to avoid tight coupling
        try:
            import numpy as np
        except ImportError:
            np = None
        self.np = np

    def _create_array_wrapper(self, array):
        """Create and returns a numpy array wrapper from a numpy array."""
        order = 'F' if (array.flags.f_contiguous and
                        not array.flags.c_contiguous) else 'C'
        allow_mmap = not self.buffered and not array.dtype.hasobject

        kwargs = {}
        try:
            self.file_handle.tell()
        except io.UnsupportedOperation:
            kwargs = {'numpy_array_alignment_bytes': None}

        wrapper = NumpyArrayWrapper(type(array),
                                    array.shape, order, array.dtype,
                                    allow_mmap=allow_mmap,
                                    **kwargs)

        return wrapper

    def save(self, obj):
        """Subclass the Pickler `save` method.

        This is a total abuse of the Pickler class in order to use the numpy
        persistence function `save` instead of the default pickle
        implementation. The numpy array is replaced by a custom wrapper in the
        pickle persistence stack and the serialized array is written right
        after in the file. Warning: the file produced does not follow the
        pickle format. As such it can not be read with `pickle.load`.
        """
        if self.np is not None and type(obj) in (self.np.ndarray,
                                                 self.np.matrix,
                                                 self.np.memmap):
            if type(obj) is self.np.memmap:
                # Pickling doesn't work with memmapped arrays
                obj = self.np.asanyarray(obj)

            # The array wrapper is pickled instead of the real array.
            wrapper = self._create_array_wrapper(obj)
            Pickler.save(self, wrapper)

            # A framer was introduced with pickle protocol 4 and we want to
            # ensure the wrapper object is written before the numpy array
            # buffer in the pickle file.
            # See https://www.python.org/dev/peps/pep-3154/#framing to get
            # more information on the framer behavior.
            if self.proto >= 4:
                self.framer.commit_frame(force=True)

            # And then array bytes are written right after the wrapper.
            wrapper.write_array(obj, self)
            return

        return Pickler.save(self, obj)

class NumpyUnpickler(Unpickler):
    """A subclass of the Unpickler to unpickle our numpy pickles.

    Attributes
    ----------
    mmap_mode: str
        The memorymap mode to use for reading numpy arrays.
    file_handle: file_like
        File object to unpickle from.
    filename: str
        Name of the file to unpickle from. It should correspond to file_handle.
        This parameter is required when using mmap_mode.
    np: module
        Reference to numpy module if numpy is installed else None.
    allow_pickle: bool
        Determine if pickle.load is allowed on the wrapped array.
        Default: False.

    """

    dispatch = Unpickler.dispatch.copy()

    def __init__(self, filename, file_handle, mmap_mode=None,
                 allow_pickle=False):
        # The next line is for backward compatibility with pickle generated
        # with joblib versions less than 0.10.
        self._dirname = os.path.dirname(filename)

        self.mmap_mode = mmap_mode
        self.allow_pickle = allow_pickle
        self.file_handle = file_handle
        # filename is required for numpy mmap mode.
        self.filename = filename
        self.compat_mode = False
        Unpickler.__init__(self, self.file_handle)
        try:
            import numpy as np
        except ImportError:
            np = None
        self.np = np

    def load_build(self):
        """Called to set the state of a newly created object.

        We capture it to replace our place-holder objects, NDArrayWrapper or
        NumpyArrayWrapper, by the array we are interested in. We
        replace them directly in the stack of pickler.
        NDArrayWrapper is used for backward compatibility with joblib <= 0.9.
        """
        Unpickler.load_build(self)

        # For backward compatibility, we support NDArrayWrapper objects.
        if isinstance(self.stack[-1], (NDArrayWrapper, NumpyArrayWrapper)):
            if self.np is None:
                raise ImportError("Trying to unpickle an ndarray, "
                                  "but numpy didn't import correctly")
            array_wrapper = self.stack.pop()
            # If any NDArrayWrapper is found, we switch to compatibility mode,
            # this will be used to raise a DeprecationWarning to the user at
            # the end of the unpickling.
            if isinstance(array_wrapper, NDArrayWrapper):
                self.compat_mode = True
                data = array_wrapper.read(self)
            else:
                data = array_wrapper.read(self, allow_pickle=self.allow_pickle)

            self.stack.append(data)

    # Be careful to register our new method.
    dispatch[pickle.BUILD[0]] = load_build

###############################################################################
# Utility functions

def dump(value, filename, compress=0, protocol=None, cache_size=None):
    """Persist an arbitrary Python object into one file.

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

    Parameters
    ----------
    value: any Python object
        The object to store to disk.
    filename: str, pathlib.Path, or file object.
        The file object or path of the file in which it is to be stored.
        The compression method corresponding to one of the supported filename
        extensions ('.z', '.gz', '.bz2', '.xz' or '.lzma') will be used
        automatically.
    compress: int from 0 to 9 or bool or 2-tuple, optional
        Optional compression level for the data. 0 or False is no compression.
        Higher value means more compression, but also slower read and
        write times. Using a value of 3 is often a good compromise.
        See the notes for more details.
        If compress is True, the compression level used is 3.
        If compress is a 2-tuple, the first element must correspond to a string
        between supported compressors (e.g 'zlib', 'gzip', 'bz2', 'lzma'
        'xz'), the second element must be an integer from 0 to 9, corresponding
        to the compression level.
    protocol: int, optional
        Pickle protocol, see pickle.dump documentation for more details.
    cache_size: positive int, optional
        This option is deprecated in 0.10 and has no effect.

    Returns
    -------
    filenames: list of strings
        The list of file names in which the data is stored. If
        compress is false, each array is stored in a different file.

    See Also
    --------
    joblib.load : corresponding loader

    Notes
    -----
    Memmapping on load cannot be used for compressed files. Thus
    using compression can significantly slow down loading. In
    addition, compressed files take up extra memory during
    dump and load.

    """

    if Path is not None and isinstance(filename, Path):
        filename = str(filename)

    is_filename = isinstance(filename, str)
    is_fileobj = hasattr(filename, "write")

    compress_method = 'zlib'  # zlib is the default compression method.
    if compress is True:
        # By default, if compress is enabled, we want the default compress
        # level of the compressor.
        compress_level = None
    elif isinstance(compress, tuple):
        # a 2-tuple was set in compress
        if len(compress) != 2:
            raise ValueError(
                'Compress argument tuple should contain exactly 2 elements: '
                '(compress method, compress level), you passed {}'
                .format(compress))
        compress_method, compress_level = compress
    elif isinstance(compress, str):
        compress_method = compress
        compress_level = None  # Use default compress level
        compress = (compress_method, compress_level)
    else:
        compress_level = compress

    if compress_method == 'lz4' and lz4 is None:
        raise ValueError(LZ4_NOT_INSTALLED_ERROR)

    if (compress_level is not None and
            compress_level is not False and
            compress_level not in range(10)):
        # Raising an error if a non valid compress level is given.
        raise ValueError(
            'Non valid compress level given: "{}". Possible values are '
            '{}.'.format(compress_level, list(range(10))))

    if compress_method not in _COMPRESSORS:
        # Raising an error if an unsupported compression method is given.
        raise ValueError(
            'Non valid compression method given: "{}". Possible values are '
            '{}.'.format(compress_method, _COMPRESSORS))

    if not is_filename and not is_fileobj:
        # People keep inverting arguments, and the resulting error is
        # incomprehensible
        raise ValueError(
            'Second argument should be a filename or a file-like object, '
            '%s (type %s) was given.'
            % (filename, type(filename))
        )

    if is_filename and not isinstance(compress, tuple):
        # In case no explicit compression was requested using both compression
        # method and level in a tuple and the filename has an explicit
        # extension, we select the corresponding compressor.

        # unset the variable to be sure no compression level is set afterwards.
        compress_method = None
        for name, compressor in _COMPRESSORS.items():
            if filename.endswith(compressor.extension):
                compress_method = name

        if compress_method in _COMPRESSORS and compress_level == 0:
            # we choose the default compress_level in case it was not given
            # as an argument (using compress).
            compress_level = None

    if cache_size is not None:
        # Cache size is deprecated starting from version 0.10
        warnings.warn("Please do not set 'cache_size' in joblib.dump, "
                      "this parameter has no effect and will be removed. "
                      "You used 'cache_size={}'".format(cache_size),
                      DeprecationWarning, stacklevel=2)

    if compress_level != 0:
        with _write_fileobject(filename, compress=(compress_method,
                                                   compress_level)) as f:
            NumpyPickler(f, protocol=protocol).dump(value)
    elif is_filename:
        with open(filename, 'wb') as f:
            NumpyPickler(f, protocol=protocol).dump(value)
    else:
        NumpyPickler(filename, protocol=protocol).dump(value)

    # If the target container is a file object, nothing is returned.
    if is_fileobj:
        return

    # For compatibility, the list of created filenames (e.g with one element
    # after 0.10.0) is returned by default.
    return [filename]

def _unpickle(fobj, filename="", mmap_mode=None, allow_pickle=False):
    """Internal unpickling function."""
    # We are careful to open the file handle early and keep it open to
    # avoid race-conditions on renames.
    # That said, if data is stored in companion files, which can be
    # the case with the old persistence format, moving the directory
    # will create a race when joblib tries to access the companion
    # files.
    unpickler = NumpyUnpickler(filename, fobj, mmap_mode=mmap_mode,
                               allow_pickle=allow_pickle)
    obj = None
    try:
        obj = unpickler.load()
        if unpickler.compat_mode:
            warnings.warn("The file '%s' has been generated with a "
                          "joblib version less than 0.10. "
                          "Please regenerate this pickle file."
                          % filename,
                          DeprecationWarning, stacklevel=3)
    except UnicodeDecodeError as exc:
        # More user-friendly error message
        new_exc = ValueError(
            'You may be trying to read with '
            'python 3 a joblib pickle generated with python 2. '
            'This feature is not supported by joblib.')
        new_exc.__cause__ = exc
        raise new_exc
    return obj

def load_temporary_memmap(filename, mmap_mode, unlink_on_gc_collect):
    from ._memmapping_reducer import JOBLIB_MMAPS, add_maybe_unlink_finalizer
    obj = load(filename, mmap_mode)
    JOBLIB_MMAPS.add(obj.filename)
    if unlink_on_gc_collect:
        add_maybe_unlink_finalizer(obj)
    return obj

def load(filename, mmap_mode=None, allow_pickle=False):
    """Reconstruct a Python object from a file persisted with joblib.dump.

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

    WARNING: joblib.load relies on the pickle module and can therefore
    execute arbitrary Python code. It should therefore never be used
    to load files from untrusted sources.

    Parameters
    ----------
    filename: str, pathlib.Path, or file object.
        The file object or path of the file from which to load the object
    mmap_mode: {None, 'r+', 'r', 'w+', 'c'}, optional
        If not None, the arrays are memory-mapped from the disk. This
        mode has no effect for compressed files. Note that in this
        case the reconstructed object might no longer match exactly
        the originally pickled object.
    allow_pickle: bool
        Determine if pickle.load is allowed on the wrapped array.
        Default: False.

    Returns
    -------
    result: any Python object
        The object stored in the file.

    See Also
    --------
    joblib.dump : function to save an object

    Notes
    -----

    This function can load numpy array files saved separately during the
    dump. If the mmap_mode argument is given, it is passed to np.load and
    arrays are loaded as memmaps. As a consequence, the reconstructed
    object might not match the original pickled object. Note that if the
    file was saved with compression, the arrays cannot be memmapped.
    """
    if Path is not None and isinstance(filename, Path):
        filename = str(filename)

    if hasattr(filename, "read"):
        fobj = filename
        filename = getattr(fobj, 'name', '')
        with _read_fileobject(fobj, filename, mmap_mode) as fobj:
            obj = _unpickle(fobj, allow_pickle=allow_pickle)
    else:
        with open(filename, 'rb') as f:
            with _read_fileobject(f, filename, mmap_mode) as fobj:
                if isinstance(fobj, str):
                    # if the returned file object is a string, this means we
                    # try to load a pickle file generated with an version of
                    # Joblib so we load it with joblib compatibility function.
                    return load_compatibility(fobj)

                obj = _unpickle(fobj, filename, mmap_mode,
                                allow_pickle)
    return obj

"""Test the numpy pickler as a replacement of the standard pickler."""

import copy
import os
import random
import re
import io
import sys
import warnings
import gzip
import zlib
import bz2
import pickle
import socket
from contextlib import closing
import mmap
from pathlib import Path

try:
    import lzma
except ImportError:
    lzma = None

import pytest

from joblib.test.common import np, with_numpy, with_lz4, without_lz4
from joblib.test.common import with_memory_profiler, memory_used
from joblib.testing import parametrize, raises, warns

# numpy_pickle is not a drop-in replacement of pickle, as it takes
# filenames instead of open files as arguments.
from joblib import numpy_pickle, register_compressor
from joblib.test import data

from joblib.numpy_pickle_utils import _IO_BUFFER_SIZE
from joblib.numpy_pickle_utils import _detect_compressor
from joblib.numpy_pickle_utils import _is_numpy_array_byte_order_mismatch
from joblib.numpy_pickle_utils import _ensure_native_byte_order
from joblib.compressor import (_COMPRESSORS, _LZ4_PREFIX, CompressorWrapper,
                               LZ4_NOT_INSTALLED_ERROR, BinaryZlibFile)

###############################################################################
# Define a list of standard types.
# Borrowed from dill, initial author: Micheal McKerns:
# http://dev.danse.us/trac/pathos/browser/dill/dill_test2.py

typelist = []

# testing types
_none = None
typelist.append(_none)
_type = type
typelist.append(_type)
_bool = bool(1)
typelist.append(_bool)
_int = int(1)
typelist.append(_int)
_float = float(1)
typelist.append(_float)
_complex = complex(1)
typelist.append(_complex)
_string = str(1)
typelist.append(_string)
_tuple = ()
typelist.append(_tuple)
_list = []
typelist.append(_list)
_dict = {}
typelist.append(_dict)
_builtin = len
typelist.append(_builtin)

def _function(x):
    yield x

class _class:
    def _method(self):
        pass

class _newclass(object):
    def _method(self):
        pass

typelist.append(_function)
typelist.append(_class)
typelist.append(_newclass)  # <type 'type'>
_instance = _class()
typelist.append(_instance)
_object = _newclass()
typelist.append(_object)  # <type 'class'>

###############################################################################
# Tests

@parametrize('compress', [0, 1])
@parametrize('member', typelist)
def test_standard_types(tmpdir, compress, member):
    # Test pickling and saving with standard types.
    filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump(member, filename, compress=compress)
    _member = numpy_pickle.load(filename)
    # We compare the pickled instance to the reloaded one only if it
    # can be compared to a copied one
    if member == copy.deepcopy(member):
        assert member == _member

def test_value_error():
    # Test inverting the input arguments to dump
    with raises(ValueError):
        numpy_pickle.dump('foo', dict())

@parametrize('wrong_compress', [-1, 10, dict()])
def test_compress_level_error(wrong_compress):
    # Verify that passing an invalid compress argument raises an error.
    exception_msg = ('Non valid compress level given: '
                     '"{0}"'.format(wrong_compress))
    with raises(ValueError) as excinfo:
        numpy_pickle.dump('dummy', 'foo', compress=wrong_compress)
    excinfo.match(exception_msg)

@with_numpy
@parametrize('compress', [False, True, 0, 3, 'zlib'])
def test_numpy_persistence(tmpdir, compress):
    filename = tmpdir.join('test.pkl').strpath
    rnd = np.random.RandomState(0)
    a = rnd.random_sample((10, 2))
    # We use 'a.T' to have a non C-contiguous array.
    for index, obj in enumerate(((a,), (a.T,), (a, a), [a, a, a])):
        filenames = numpy_pickle.dump(obj, filename, compress=compress)

        # All is cached in one file
        assert len(filenames) == 1
        # Check that only one file was created
        assert filenames[0] == filename
        # Check that this file does exist
        assert os.path.exists(filenames[0])

        # Unpickle the object
        obj_ = numpy_pickle.load(filename)
        # Check that the items are indeed arrays
        for item in obj_:
            assert isinstance(item, np.ndarray)
        # And finally, check that all the values are equal.
        np.testing.assert_array_equal(np.array(obj), np.array(obj_))

    # Now test with an array subclass
    obj = np.memmap(filename + 'mmap', mode='w+', shape=4, dtype=np.float64)
    filenames = numpy_pickle.dump(obj, filename, compress=compress)
    # All is cached in one file
    assert len(filenames) == 1

    obj_ = numpy_pickle.load(filename)
    if (type(obj) is not np.memmap and
            hasattr(obj, '__array_prepare__')):
        # We don't reconstruct memmaps
        assert isinstance(obj_, type(obj))

    np.testing.assert_array_equal(obj_, obj)

    # Test with an object containing multiple numpy arrays
    obj = ComplexTestObject()
    filenames = numpy_pickle.dump(obj, filename, compress=compress)
    # All is cached in one file
    assert len(filenames) == 1

    obj_loaded = numpy_pickle.load(filename, allow_pickle=True)
    assert isinstance(obj_loaded, type(obj))
    np.testing.assert_array_equal(obj_loaded.array_float, obj.array_float)
    np.testing.assert_array_equal(obj_loaded.array_int, obj.array_int)
    np.testing.assert_array_equal(obj_loaded.array_obj, obj.array_obj)

@with_numpy
def test_numpy_persistence_bufferred_array_compression(tmpdir):
    big_array = np.ones((_IO_BUFFER_SIZE + 100), dtype=np.uint8)
    filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump(big_array, filename, compress=True)
    arr_reloaded = numpy_pickle.load(filename)

    np.testing.assert_array_equal(big_array, arr_reloaded)

@with_numpy
def test_memmap_persistence(tmpdir):
    rnd = np.random.RandomState(0)
    a = rnd.random_sample(10)
    filename = tmpdir.join('test1.pkl').strpath
    numpy_pickle.dump(a, filename)
    b = numpy_pickle.load(filename, mmap_mode='r')

    assert isinstance(b, np.memmap)

    # Test with an object containing multiple numpy arrays
    filename = tmpdir.join('test2.pkl').strpath
    obj = ComplexTestObject()
    numpy_pickle.dump(obj, filename)
    obj_loaded = numpy_pickle.load(filename, mmap_mode='r',
                                   allow_pickle=True)
    assert isinstance(obj_loaded, type(obj))
    assert isinstance(obj_loaded.array_float, np.memmap)
    assert not obj_loaded.array_float.flags.writeable
    assert isinstance(obj_loaded.array_int, np.memmap)
    assert not obj_loaded.array_int.flags.writeable
    # Memory map not allowed for numpy object arrays
    assert not isinstance(obj_loaded.array_obj, np.memmap)
    np.testing.assert_array_equal(obj_loaded.array_float,
                                  obj.array_float)
    np.testing.assert_array_equal(obj_loaded.array_int,
                                  obj.array_int)
    np.testing.assert_array_equal(obj_loaded.array_obj,
                                  obj.array_obj)

    # Test we can write in memmapped arrays
    obj_loaded = numpy_pickle.load(filename, mmap_mode='r+',
                                   allow_pickle=True)
    assert obj_loaded.array_float.flags.writeable
    obj_loaded.array_float[0:10] = 10.0
    assert obj_loaded.array_int.flags.writeable
    obj_loaded.array_int[0:10] = 10

    obj_reloaded = numpy_pickle.load(filename, mmap_mode='r',
                                     allow_pickle=True)
    np.testing.assert_array_equal(obj_reloaded.array_float,
                                  obj_loaded.array_float)
    np.testing.assert_array_equal(obj_reloaded.array_int,
                                  obj_loaded.array_int)

    # Test w+ mode is caught and the mode has switched to r+
    numpy_pickle.load(filename, mmap_mode='w+', allow_pickle=True)
    assert obj_loaded.array_int.flags.writeable
    assert obj_loaded.array_int.mode == 'r+'
    assert obj_loaded.array_float.flags.writeable
    assert obj_loaded.array_float.mode == 'r+'

@with_numpy
def test_memmap_persistence_mixed_dtypes(tmpdir):
    # loading datastructures that have sub-arrays with dtype=object
    # should not prevent memmapping on fixed size dtype sub-arrays.
    rnd = np.random.RandomState(0)
    a = rnd.random_sample(10)
    b = np.array([1, 'b'], dtype=object)
    construct = (a, b)
    filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump(construct, filename)
    a_clone, b_clone = numpy_pickle.load(filename, mmap_mode='r',
                                         allow_pickle=True)

    # the floating point array has been memory mapped
    assert isinstance(a_clone, np.memmap)

    # the object-dtype array has been loaded in memory
    assert not isinstance(b_clone, np.memmap)

@with_numpy
def test_masked_array_persistence(tmpdir):
    # The special-case picker fails, because saving masked_array
    # not implemented, but it just delegates to the standard pickler.
    rnd = np.random.RandomState(0)
    a = rnd.random_sample(10)
    a = np.ma.masked_greater(a, 0.5)
    filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump(a, filename)
    b = numpy_pickle.load(filename, mmap_mode='r')
    assert isinstance(b, np.ma.masked_array)

@with_numpy
def test_compress_mmap_mode_warning(tmpdir):
    # Test the warning in case of compress + mmap_mode
    rnd = np.random.RandomState(0)
    a = rnd.random_sample(10)
    this_filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump(a, this_filename, compress=1)
    with warns(UserWarning) as warninfo:
        numpy_pickle.load(this_filename, mmap_mode='r+')
    debug_msg = "\n".join([str(w) for w in warninfo])
    warninfo = [w.message for w in warninfo]
    assert len(warninfo) == 1, debug_msg
    assert (
        str(warninfo[0]) ==
        'mmap_mode "r+" is not compatible with compressed '
        f'file {this_filename}. "r+" flag will be ignored.'
    )

@with_numpy
@parametrize('cache_size', [None, 0, 10])
def test_cache_size_warning(tmpdir, cache_size):
    # Check deprecation warning raised when cache size is not None
    filename = tmpdir.join('test.pkl').strpath
    rnd = np.random.RandomState(0)
    a = rnd.random_sample((10, 2))

    warnings.simplefilter("always")
    with warnings.catch_warnings(record=True) as warninfo:
        numpy_pickle.dump(a, filename, cache_size=cache_size)
    expected_nb_warnings = 1 if cache_size is not None else 0
    assert len(warninfo) == expected_nb_warnings
    for w in warninfo:
        assert w.category == DeprecationWarning
        assert (str(w.message) ==
                "Please do not set 'cache_size' in joblib.dump, this "
                "parameter has no effect and will be removed. You "
                "used 'cache_size={0}'".format(cache_size))

@with_numpy
@with_memory_profiler
@parametrize('compress', [True, False])
def test_memory_usage(tmpdir, compress):
    # Verify memory stays within expected bounds.
    filename = tmpdir.join('test.pkl').strpath
    small_array = np.ones((10, 10))
    big_array = np.ones(shape=100 * int(1e6), dtype=np.uint8)

    for obj in (small_array, big_array):
        size = obj.nbytes / 1e6
        obj_filename = filename + str(np.random.randint(0, 1000))
        mem_used = memory_used(numpy_pickle.dump,
                               obj, obj_filename, compress=compress)

        # The memory used to dump the object shouldn't exceed the buffer
        # size used to write array chunks (16MB).
        write_buf_size = _IO_BUFFER_SIZE + 16 * 1024 ** 2 / 1e6
        assert mem_used <= write_buf_size

        mem_used = memory_used(numpy_pickle.load, obj_filename)
        # memory used should be less than array size + buffer size used to
        # read the array chunk by chunk.
        read_buf_size = 32 + _IO_BUFFER_SIZE  # MiB
        assert mem_used < size + read_buf_size

@with_numpy
def test_compressed_pickle_dump_and_load(tmpdir):
    expected_list = [np.arange(5, dtype=np.dtype('<i8')),
                     np.arange(5, dtype=np.dtype('>i8')),
                     np.arange(5, dtype=np.dtype('<f8')),
                     np.arange(5, dtype=np.dtype('>f8')),
                     np.array([1, 'abc', {'a': 1, 'b': 2}], dtype='O'),
                     np.arange(256, dtype=np.uint8).tobytes(),
                     u"C'est l'\xe9t\xe9 !"]

    fname = tmpdir.join('temp.pkl.gz').strpath

    dumped_filenames = numpy_pickle.dump(expected_list, fname, compress=1)
    assert len(dumped_filenames) == 1
    result_list = numpy_pickle.load(fname, allow_pickle=True)
    for result, expected in zip(result_list, expected_list):
        if isinstance(expected, np.ndarray):
            expected = _ensure_native_byte_order(expected)
            assert result.dtype == expected.dtype
            np.testing.assert_equal(result, expected)
        else:
            assert result == expected

def _check_pickle(filename, expected_list, mmap_mode=None):
    """Helper function to test joblib pickle content.

    Note: currently only pickles containing an iterable are supported
    by this function.
    """
    version_match = re.match(r'.+py(\d)(\d).+', filename)
    py_version_used_for_writing = int(version_match.group(1))

    py_version_to_default_pickle_protocol = {2: 2, 3: 3}
    pickle_reading_protocol = py_version_to_default_pickle_protocol.get(3, 4)
    pickle_writing_protocol = py_version_to_default_pickle_protocol.get(
        py_version_used_for_writing, 4)
    if pickle_reading_protocol >= pickle_writing_protocol:
        try:
            with warnings.catch_warnings(record=True) as warninfo:
                warnings.simplefilter('always')
                warnings.filterwarnings(
                    'ignore', module='numpy',
                    message='The compiler package is deprecated')
                result_list = numpy_pickle.load(filename,
                                                mmap_mode=mmap_mode,
                                                allow_pickle=True)
            filename_base = os.path.basename(filename)
            expected_nb_deprecation_warnings = 1 if (
                "_0.9" in filename_base or "_0.8.4" in filename_base) else 0

            expected_nb_user_warnings = 3 if (
                re.search("_0.1.+.pkl$", filename_base) and
                mmap_mode is not None) else 0
            expected_nb_warnings = \
                expected_nb_deprecation_warnings + expected_nb_user_warnings
            assert len(warninfo) == expected_nb_warnings

            deprecation_warnings = [
                w for w in warninfo if issubclass(
                    w.category, DeprecationWarning)]
            user_warnings = [
                w for w in warninfo if issubclass(
                    w.category, UserWarning)]
            for w in deprecation_warnings:
                assert (str(w.message) ==
                        "The file '{0}' has been generated with a joblib "
                        "version less than 0.10. Please regenerate this "
                        "pickle file.".format(filename))

            for w in user_warnings:
                escaped_filename = re.escape(filename)
                assert re.search(
                    f"memmapped.+{escaped_filename}.+segmentation fault",
                    str(w.message))

            for result, expected in zip(result_list, expected_list):
                if isinstance(expected, np.ndarray):
                    expected = _ensure_native_byte_order(expected)
                    assert result.dtype == expected.dtype
                    np.testing.assert_equal(result, expected)
                else:
                    assert result == expected
        except Exception as exc:
            # When trying to read with python 3 a pickle generated
            # with python 2 we expect a user-friendly error
            if py_version_used_for_writing == 2:
                assert isinstance(exc, ValueError)
                message = ('You may be trying to read with '
                           'python 3 a joblib pickle generated with python 2.')
                assert message in str(exc)
            elif filename.endswith('.lz4') and with_lz4.args[0]:
                assert isinstance(exc, ValueError)
                assert LZ4_NOT_INSTALLED_ERROR in str(exc)
            else:
                raise
    else:
        # Pickle protocol used for writing is too high. We expect a
        # "unsupported pickle protocol" error message
        try:
            numpy_pickle.load(filename)
            raise AssertionError('Numpy pickle loading should '
                                 'have raised a ValueError exception')
        except ValueError as e:
            message = 'unsupported pickle protocol: {0}'.format(
                pickle_writing_protocol)
            assert message in str(e.args)

@with_numpy
def test_joblib_pickle_across_python_versions():
    # We need to be specific about dtypes in particular endianness
    # because the pickles can be generated on one architecture and
    # the tests run on another one. See
    # https://github.com/joblib/joblib/issues/279.
    expected_list = [np.arange(5, dtype=np.dtype('<i8')),
                     np.arange(5, dtype=np.dtype('<f8')),
                     np.array([1, 'abc', {'a': 1, 'b': 2}], dtype='O'),
                     np.arange(256, dtype=np.uint8).tobytes(),
                     # np.matrix is a subclass of np.ndarray, here we want
                     # to verify this type of object is correctly unpickled
                     # among versions.
                     np.matrix([0, 1, 2], dtype=np.dtype('<i8')),
                     u"C'est l'\xe9t\xe9 !"]

    # Testing all the compressed and non compressed
    # pickles in joblib/test/data. These pickles were generated by
    # the joblib/test/data/create_numpy_pickle.py script for the
    # relevant python, joblib and numpy versions.
    test_data_dir = os.path.dirname(os.path.abspath(data.__file__))

    pickle_extensions = ('.pkl', '.gz', '.gzip', '.bz2', 'lz4')
    if lzma is not None:
        pickle_extensions += ('.xz', '.lzma')
    pickle_filenames = [os.path.join(test_data_dir, fn)
                        for fn in os.listdir(test_data_dir)
                        if any(fn.endswith(ext) for ext in pickle_extensions)]

    for fname in pickle_filenames:
        _check_pickle(fname, expected_list)

@with_numpy
def test_joblib_pickle_across_python_versions_with_mmap():
    expected_list = [np.arange(5, dtype=np.dtype('<i8')),
                     np.arange(5, dtype=np.dtype('<f8')),
                     np.array([1, 'abc', {'a': 1, 'b': 2}], dtype='O'),
                     np.arange(256, dtype=np.uint8).tobytes(),
                     # np.matrix is a subclass of np.ndarray, here we want
                     # to verify this type of object is correctly unpickled
                     # among versions.
                     np.matrix([0, 1, 2], dtype=np.dtype('<i8')),
                     u"C'est l'\xe9t\xe9 !"]

    test_data_dir = os.path.dirname(os.path.abspath(data.__file__))

    pickle_filenames = [
        os.path.join(test_data_dir, fn)
        for fn in os.listdir(test_data_dir) if fn.endswith('.pkl')]
    for fname in pickle_filenames:
        _check_pickle(fname, expected_list, mmap_mode='r')

@with_numpy
def test_numpy_array_byte_order_mismatch_detection():
    # List of numpy arrays with big endian byteorder.
    be_arrays = [np.array([(1, 2.0), (3, 4.0)],
                          dtype=[('', '>i8'), ('', '>f8')]),
                 np.arange(3, dtype=np.dtype('>i8')),
                 np.arange(3, dtype=np.dtype('>f8'))]

    # Verify the byteorder mismatch is correctly detected.
    for array in be_arrays:
        if sys.byteorder == 'big':
            assert not _is_numpy_array_byte_order_mismatch(array)
        else:
            assert _is_numpy_array_byte_order_mismatch(array)
        converted = _ensure_native_byte_order(array)
        if converted.dtype.fields:
            for f in converted.dtype.fields.values():
                f[0].byteorder == '='
        else:
            assert converted.dtype.byteorder == "="

    # List of numpy arrays with little endian byteorder.
    le_arrays = [np.array([(1, 2.0), (3, 4.0)],
                          dtype=[('', '<i8'), ('', '<f8')]),
                 np.arange(3, dtype=np.dtype('<i8')),
                 np.arange(3, dtype=np.dtype('<f8'))]

    # Verify the byteorder mismatch is correctly detected.
    for array in le_arrays:
        if sys.byteorder == 'little':
            assert not _is_numpy_array_byte_order_mismatch(array)
        else:
            assert _is_numpy_array_byte_order_mismatch(array)
        converted = _ensure_native_byte_order(array)
        if converted.dtype.fields:
            for f in converted.dtype.fields.values():
                f[0].byteorder == '='
        else:
            assert converted.dtype.byteorder == "="

@parametrize('compress_tuple', [('zlib', 3), ('gzip', 3)])
def test_compress_tuple_argument(tmpdir, compress_tuple):
    # Verify the tuple is correctly taken into account.
    filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump("dummy", filename,
                      compress=compress_tuple)
    # Verify the file contains the right magic number
    with open(filename, 'rb') as f:
        assert _detect_compressor(f) == compress_tuple[0]

@parametrize('compress_tuple,message',
             [(('zlib', 3, 'extra'),        # wrong compress tuple
               'Compress argument tuple should contain exactly 2 elements'),
              (('wrong', 3),                # wrong compress method
               'Non valid compression method given: "{}"'.format('wrong')),
              (('zlib', 'wrong'),           # wrong compress level
               'Non valid compress level given: "{}"'.format('wrong'))])
def test_compress_tuple_argument_exception(tmpdir, compress_tuple, message):
    filename = tmpdir.join('test.pkl').strpath
    # Verify setting a wrong compress tuple raises a ValueError.
    with raises(ValueError) as excinfo:
        numpy_pickle.dump('dummy', filename, compress=compress_tuple)
    excinfo.match(message)

@parametrize('compress_string', ['zlib', 'gzip'])
def test_compress_string_argument(tmpdir, compress_string):
    # Verify the string is correctly taken into account.
    filename = tmpdir.join('test.pkl').strpath
    numpy_pickle.dump("dummy", filename,
                      compress=compress_string)
    # Verify the file contains the right magic number
    with open(filename, 'rb') as f:
        assert _detect_compressor(f) == compress_string

@with_numpy
@parametrize('compress', [1, 3, 6])
@parametrize('cmethod', _COMPRESSORS)
def test_joblib_compression_formats(tmpdir, compress, cmethod):
    filename = tmpdir.join('test.pkl').strpath
    objects = (np.ones(shape=(100, 100), dtype='f8'),
               range(10),
               {'a': 1, 2: 'b'}, [], (), {}, 0, 1.0)

    if cmethod in ("lzma", "xz") and lzma is None:
        pytest.skip("lzma is support not available")

    elif cmethod == 'lz4' and with_lz4.args[0]:
        # Skip the test if lz4 is not installed. We here use the with_lz4
        # skipif fixture whose argument is True when lz4 is not installed
        pytest.skip("lz4 is not installed.")

    dump_filename = filename + "." + cmethod
    for obj in objects:
        numpy_pickle.dump(obj, dump_filename, compress=(cmethod, compress))
        # Verify the file contains the right magic number
        with open(dump_filename, 'rb') as f:
            assert _detect_compressor(f) == cmethod
        # Verify the reloaded object is correct
        obj_reloaded = numpy_pickle.load(dump_filename)
        assert isinstance(obj_reloaded, type(obj))
        if isinstance(obj, np.ndarray):
            np.testing.assert_array_equal(obj_reloaded, obj)
        else:
            assert obj_reloaded == obj

def _gzip_file_decompress(source_filename, target_filename):
    """Decompress a gzip file."""
    with closing(gzip.GzipFile(source_filename, "rb")) as fo:
        buf = fo.read()

    with open(target_filename, "wb") as fo:
        fo.write(buf)

def _zlib_file_decompress(source_filename, target_filename):
    """Decompress a zlib file."""
    with open(source_filename, 'rb') as fo:
        buf = zlib.decompress(fo.read())

    with open(target_filename, 'wb') as fo:
        fo.write(buf)

@parametrize('extension,decompress',
             [('.z', _zlib_file_decompress),
              ('.gz', _gzip_file_decompress)])
def test_load_externally_decompressed_files(tmpdir, extension, decompress):
    # Test that BinaryZlibFile generates valid gzip and zlib compressed files.
    obj = "a string to persist"
    filename_raw = tmpdir.join('test.pkl').strpath

    filename_compressed = filename_raw + extension
    # Use automatic extension detection to compress with the right method.
    numpy_pickle.dump(obj, filename_compressed)

    # Decompress with the corresponding method
    decompress(filename_compressed, filename_raw)

    # Test that the uncompressed pickle can be loaded and
    # that the result is correct.
    obj_reloaded = numpy_pickle.load(filename_raw)
    assert obj == obj_reloaded

@parametrize('extension,cmethod',
             # valid compressor extensions
             [('.z', 'zlib'),
              ('.gz', 'gzip'),
              ('.bz2', 'bz2'),
              ('.lzma', 'lzma'),
              ('.xz', 'xz'),
              # invalid compressor extensions
              ('.pkl', 'not-compressed'),
              ('', 'not-compressed')])
def test_compression_using_file_extension(tmpdir, extension, cmethod):
    if cmethod in ("lzma", "xz") and lzma is None:
        pytest.skip("lzma is missing")
    # test that compression method corresponds to the given filename extension.
    filename = tmpdir.join('test.pkl').strpath
    obj = "object to dump"

    dump_fname = filename + extension
    numpy_pickle.dump(obj, dump_fname)
    # Verify the file contains the right magic number
    with open(dump_fname, 'rb') as f:
        assert _detect_compressor(f) == cmethod
    # Verify the reloaded object is correct
    obj_reloaded = numpy_pickle.load(dump_fname)
    assert isinstance(obj_reloaded, type(obj))
    assert obj_reloaded == obj

@with_numpy
def test_file_handle_persistence(tmpdir):
    objs = [np.random.random((10, 10)), "some data"]
    fobjs = [bz2.BZ2File, gzip.GzipFile]
    if lzma is not None:
        fobjs += [lzma.LZMAFile]
    filename = tmpdir.join('test.pkl').strpath

    for obj in objs:
        for fobj in fobjs:
            with fobj(filename, 'wb') as f:
                numpy_pickle.dump(obj, f)

            # using the same decompressor prevents from internally
            # decompress again.
            with fobj(filename, 'rb') as f:
                obj_reloaded = numpy_pickle.load(f)

            # when needed, the correct decompressor should be used when
            # passing a raw file handle.
            with open(filename, 'rb') as f:
                obj_reloaded_2 = numpy_pickle.load(f)

            if isinstance(obj, np.ndarray):
                np.testing.assert_array_equal(obj_reloaded, obj)
                np.testing.assert_array_equal(obj_reloaded_2, obj)
            else:
                assert obj_reloaded == obj
                assert obj_reloaded_2 == obj

@with_numpy
def test_in_memory_persistence():
    objs = [np.random.random((10, 10)), "some data"]
    for obj in objs:
        f = io.BytesIO()
        numpy_pickle.dump(obj, f)
        obj_reloaded = numpy_pickle.load(f)
        if isinstance(obj, np.ndarray):
            np.testing.assert_array_equal(obj_reloaded, obj)
        else:
            assert obj_reloaded == obj

@with_numpy
def test_file_handle_persistence_mmap(tmpdir):
    obj = np.random.random((10, 10))
    filename = tmpdir.join('test.pkl').strpath

    with open(filename, 'wb') as f:
        numpy_pickle.dump(obj, f)

    with open(filename, 'rb') as f:
        obj_reloaded = numpy_pickle.load(f, mmap_mode='r+')

    np.testing.assert_array_equal(obj_reloaded, obj)

@with_numpy
def test_file_handle_persistence_compressed_mmap(tmpdir):
    obj = np.random.random((10, 10))
    filename = tmpdir.join('test.pkl').strpath

    with open(filename, 'wb') as f:
        numpy_pickle.dump(obj, f, compress=('gzip', 3))

    with closing(gzip.GzipFile(filename, 'rb')) as f:
        with warns(UserWarning) as warninfo:
            numpy_pickle.load(f, mmap_mode='r+')
        assert len(warninfo) == 1
        assert (str(warninfo[0].message) ==
                '"%(fileobj)r" is not a raw file, mmap_mode "%(mmap_mode)s" '
                'flag will be ignored.' % {'fileobj': f, 'mmap_mode': 'r+'})

@with_numpy
def test_file_handle_persistence_in_memory_mmap():
    obj = np.random.random((10, 10))
    buf = io.BytesIO()

    numpy_pickle.dump(obj, buf)

    with warns(UserWarning) as warninfo:
        numpy_pickle.load(buf, mmap_mode='r+')
    assert len(warninfo) == 1
    assert (str(warninfo[0].message) ==
            'In memory persistence is not compatible with mmap_mode '
            '"%(mmap_mode)s" flag passed. mmap_mode option will be '
            'ignored.' % {'mmap_mode': 'r+'})

@parametrize('data', [b'a little data as bytes.',
                      # More bytes
                      10000 * "{}".format(
                          random.randint(0, 1000) * 1000).encode('latin-1')],
             ids=["a little data as bytes.", "a large data as bytes."])
@parametrize('compress_level', [1, 3, 9])
def test_binary_zlibfile(tmpdir, data, compress_level):
    filename = tmpdir.join('test.pkl').strpath
    # Regular cases
    with open(filename, 'wb') as f:
        with BinaryZlibFile(f, 'wb',
                            compresslevel=compress_level) as fz:
            assert fz.writable()
            fz.write(data)
            assert fz.fileno() == f.fileno()
            with raises(io.UnsupportedOperation):
                fz._check_can_read()

            with raises(io.UnsupportedOperation):
                fz._check_can_seek()
        assert fz.closed
        with raises(ValueError):
            fz._check_not_closed()

    with open(filename, 'rb') as f:
        with BinaryZlibFile(f) as fz:
            assert fz.readable()
            assert fz.seekable()
            assert fz.fileno() == f.fileno()
            assert fz.read() == data
            with raises(io.UnsupportedOperation):
                fz._check_can_write()
            assert fz.seekable()
            fz.seek(0)
            assert fz.tell() == 0
        assert fz.closed

    # Test with a filename as input
    with BinaryZlibFile(filename, 'wb',
                        compresslevel=compress_level) as fz:
        assert fz.writable()
        fz.write(data)

    with BinaryZlibFile(filename, 'rb') as fz:
        assert fz.read() == data
        assert fz.seekable()

    # Test without context manager
    fz = BinaryZlibFile(filename, 'wb', compresslevel=compress_level)
    assert fz.writable()
    fz.write(data)
    fz.close()

    fz = BinaryZlibFile(filename, 'rb')
    assert fz.read() == data
    fz.close()

@parametrize('bad_value', [-1, 10, 15, 'a', (), {}])
def test_binary_zlibfile_bad_compression_levels(tmpdir, bad_value):
    filename = tmpdir.join('test.pkl').strpath
    with raises(ValueError) as excinfo:
        BinaryZlibFile(filename, 'wb', compresslevel=bad_value)
    pattern = re.escape("'compresslevel' must be an integer between 1 and 9. "
                        "You provided 'compresslevel={}'".format(bad_value))
    excinfo.match(pattern)

@parametrize('bad_mode', ['a', 'x', 'r', 'w', 1, 2])
def test_binary_zlibfile_invalid_modes(tmpdir, bad_mode):
    filename = tmpdir.join('test.pkl').strpath
    with raises(ValueError) as excinfo:
        BinaryZlibFile(filename, bad_mode)
    excinfo.match("Invalid mode")

@parametrize('bad_file', [1, (), {}])
def test_binary_zlibfile_invalid_filename_type(bad_file):
    with raises(TypeError) as excinfo:
        BinaryZlibFile(bad_file, 'rb')
    excinfo.match("filename must be a str or bytes object, or a file")

###############################################################################
# Test dumping array subclasses
if np is not None:

    class SubArray(np.ndarray):

        def __reduce__(self):
            return _load_sub_array, (np.asarray(self), )

    def _load_sub_array(arr):
        d = SubArray(arr.shape)
        d[:] = arr
        return d

    class ComplexTestObject:
        """A complex object containing numpy arrays as attributes."""

        def __init__(self):
            self.array_float = np.arange(100, dtype='float64')
            self.array_int = np.ones(100, dtype='int32')
            self.array_obj = np.array(['a', 10, 20.0], dtype='object')

@with_numpy
def test_numpy_subclass(tmpdir):
    filename = tmpdir.join('test.pkl').strpath
    a = SubArray((10,))
    numpy_pickle.dump(a, filename)
    c = numpy_pickle.load(filename)
    assert isinstance(c, SubArray)
    np.testing.assert_array_equal(c, a)

def test_pathlib(tmpdir):
    filename = tmpdir.join('test.pkl').strpath
    value = 123
    numpy_pickle.dump(value, Path(filename))
    assert numpy_pickle.load(filename) == value
    numpy_pickle.dump(value, filename)
    assert numpy_pickle.load(Path(filename)) == value

@with_numpy
def test_non_contiguous_array_pickling(tmpdir):
    filename = tmpdir.join('test.pkl').strpath

    for array in [  # Array that triggers a contiguousness issue with nditer,
                    # see https://github.com/joblib/joblib/pull/352 and see
                    # https://github.com/joblib/joblib/pull/353
                    np.asfortranarray([[1, 2], [3, 4]])[1:],
                    # Non contiguous array with works fine with nditer
                    np.ones((10, 50, 20), order='F')[:, :1, :]]:
        assert not array.flags.c_contiguous
        assert not array.flags.f_contiguous
        numpy_pickle.dump(array, filename)
        array_reloaded = numpy_pickle.load(filename)
        np.testing.assert_array_equal(array_reloaded, array)

@with_numpy
def test_pickle_highest_protocol(tmpdir):
    # ensure persistence of a numpy array is valid even when using
    # the pickle HIGHEST_PROTOCOL.
    # see https://github.com/joblib/joblib/issues/362

    filename = tmpdir.join('test.pkl').strpath
    test_array = np.zeros(10)

    numpy_pickle.dump(test_array, filename, protocol=pickle.HIGHEST_PROTOCOL)
    array_reloaded = numpy_pickle.load(filename)

    np.testing.assert_array_equal(array_reloaded, test_array)

@with_numpy
def test_pickle_in_socket():
    # test that joblib can pickle in sockets
    test_array = np.arange(10)
    _ADDR = ("localhost", 12345)
    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    listener.bind(_ADDR)
    listener.listen(1)

    with socket.create_connection(_ADDR) as client:
        server, client_addr = listener.accept()

        with server.makefile("wb") as sf:
            numpy_pickle.dump(test_array, sf)

        with client.makefile("rb") as cf:
            array_reloaded = numpy_pickle.load(cf)

        np.testing.assert_array_equal(array_reloaded, test_array)

        # Check that a byte-aligned numpy array written in a file can be send
        # over a socket and then read on the other side
        bytes_to_send = io.BytesIO()
        numpy_pickle.dump(test_array, bytes_to_send)
        server.send(bytes_to_send.getvalue())

        with client.makefile("rb") as cf:
            array_reloaded = numpy_pickle.load(cf)

        np.testing.assert_array_equal(array_reloaded, test_array)

@with_numpy
def test_load_memmap_with_big_offset(tmpdir):
    # Test that numpy memmap offset is set correctly if greater than
    # mmap.ALLOCATIONGRANULARITY, see
    # https://github.com/joblib/joblib/issues/451 and
    # https://github.com/numpy/numpy/pull/8443 for more details.
    fname = tmpdir.join('test.mmap').strpath
    size = mmap.ALLOCATIONGRANULARITY
    obj = [np.zeros(size, dtype='uint8'), np.ones(size, dtype='uint8')]
    numpy_pickle.dump(obj, fname)
    memmaps = numpy_pickle.load(fname, mmap_mode='r')
    assert isinstance(memmaps[1], np.memmap)
    assert memmaps[1].offset > size
    np.testing.assert_array_equal(obj, memmaps)

def test_register_compressor(tmpdir):
    # Check that registering compressor file works.
    compressor_name = 'test-name'
    compressor_prefix = 'test-prefix'

    class BinaryCompressorTestFile(io.BufferedIOBase):
        pass

    class BinaryCompressorTestWrapper(CompressorWrapper):

        def __init__(self):
            CompressorWrapper.__init__(self, obj=BinaryCompressorTestFile,
                                       prefix=compressor_prefix)

    register_compressor(compressor_name, BinaryCompressorTestWrapper())

    assert (_COMPRESSORS[compressor_name].fileobj_factory ==
            BinaryCompressorTestFile)
    assert _COMPRESSORS[compressor_name].prefix == compressor_prefix

    # Remove this dummy compressor file from extra compressors because other
    # tests might fail because of this
    _COMPRESSORS.pop(compressor_name)

@parametrize('invalid_name', [1, (), {}])
def test_register_compressor_invalid_name(invalid_name):
    # Test that registering an invalid compressor name is not allowed.
    with raises(ValueError) as excinfo:
        register_compressor(invalid_name, None)
    excinfo.match("Compressor name should be a string")

def test_register_compressor_invalid_fileobj():
    # Test that registering an invalid file object is not allowed.

    class InvalidFileObject():
        pass

    class InvalidFileObjectWrapper(CompressorWrapper):
        def __init__(self):
            CompressorWrapper.__init__(self, obj=InvalidFileObject,
                                       prefix=b'prefix')

    with raises(ValueError) as excinfo:
        register_compressor('invalid', InvalidFileObjectWrapper())

    excinfo.match("Compressor 'fileobj_factory' attribute should implement "
                  "the file object interface")

class AnotherZlibCompressorWrapper(CompressorWrapper):

    def __init__(self):
        CompressorWrapper.__init__(self, obj=BinaryZlibFile, prefix=b'prefix')

class StandardLibGzipCompressorWrapper(CompressorWrapper):

    def __init__(self):
        CompressorWrapper.__init__(self, obj=gzip.GzipFile, prefix=b'prefix')

def test_register_compressor_already_registered():
    # Test registration of existing compressor files.
    compressor_name = 'test-name'

    # register a test compressor
    register_compressor(compressor_name, AnotherZlibCompressorWrapper())

    with raises(ValueError) as excinfo:
        register_compressor(compressor_name,
                            StandardLibGzipCompressorWrapper())
    excinfo.match("Compressor '{}' already registered."
                  .format(compressor_name))

    register_compressor(compressor_name, StandardLibGzipCompressorWrapper(),
                        force=True)

    assert compressor_name in _COMPRESSORS
    assert _COMPRESSORS[compressor_name].fileobj_factory == gzip.GzipFile

    # Remove this dummy compressor file from extra compressors because other
    # tests might fail because of this
    _COMPRESSORS.pop(compressor_name)

@with_lz4
def test_lz4_compression(tmpdir):
    # Check that lz4 can be used when dependency is available.
    import lz4.frame
    compressor = 'lz4'
    assert compressor in _COMPRESSORS
    assert _COMPRESSORS[compressor].fileobj_factory == lz4.frame.LZ4FrameFile

    fname = tmpdir.join('test.pkl').strpath
    data = 'test data'
    numpy_pickle.dump(data, fname, compress=compressor)

    with open(fname, 'rb') as f:
        assert f.read(len(_LZ4_PREFIX)) == _LZ4_PREFIX
    assert numpy_pickle.load(fname) == data

    # Test that LZ4 is applied based on file extension
    numpy_pickle.dump(data, fname + '.lz4')
    with open(fname, 'rb') as f:
        assert f.read(len(_LZ4_PREFIX)) == _LZ4_PREFIX
    assert numpy_pickle.load(fname) == data

@without_lz4
def test_lz4_compression_without_lz4(tmpdir):
    # Check that lz4 cannot be used when dependency is not available.
    fname = tmpdir.join('test.nolz4').strpath
    data = 'test data'
    msg = LZ4_NOT_INSTALLED_ERROR
    with raises(ValueError) as excinfo:
        numpy_pickle.dump(data, fname, compress='lz4')
    excinfo.match(msg)

    with raises(ValueError) as excinfo:
        numpy_pickle.dump(data, fname + '.lz4')
    excinfo.match(msg)

protocols = [pickle.DEFAULT_PROTOCOL]
if pickle.HIGHEST_PROTOCOL != pickle.DEFAULT_PROTOCOL:
    protocols.append(pickle.HIGHEST_PROTOCOL)

@with_numpy
@parametrize('protocol', protocols)
def test_memmap_alignment_padding(tmpdir, protocol):
    # Test that memmaped arrays returned by numpy.load are correctly aligned
    fname = tmpdir.join('test.mmap').strpath

    a = np.random.randn(2)
    numpy_pickle.dump(a, fname, protocol=protocol)
    memmap = numpy_pickle.load(fname, mmap_mode='r')
    assert isinstance(memmap, np.memmap)
    np.testing.assert_array_equal(a, memmap)
    assert (
        memmap.ctypes.data % numpy_pickle.NUMPY_ARRAY_ALIGNMENT_BYTES == 0)
    assert memmap.flags.aligned

    array_list = [
        np.random.randn(2), np.random.randn(2),
        np.random.randn(2), np.random.randn(2)
    ]

    # On Windows OSError 22 if reusing the same path for memmap ...
    fname = tmpdir.join('test1.mmap').strpath
    numpy_pickle.dump(array_list, fname, protocol=protocol)
    l_reloaded = numpy_pickle.load(fname, mmap_mode='r')

    for idx, memmap in enumerate(l_reloaded):
        assert isinstance(memmap, np.memmap)
        np.testing.assert_array_equal(array_list[idx], memmap)
        assert (
            memmap.ctypes.data % numpy_pickle.NUMPY_ARRAY_ALIGNMENT_BYTES == 0)
        assert memmap.flags.aligned

    array_dict = {
        'a0': np.arange(2, dtype=np.uint8),
        'a1': np.arange(3, dtype=np.uint8),
        'a2': np.arange(5, dtype=np.uint8),
        'a3': np.arange(7, dtype=np.uint8),
        'a4': np.arange(11, dtype=np.uint8),
        'a5': np.arange(13, dtype=np.uint8),
        'a6': np.arange(17, dtype=np.uint8),
        'a7': np.arange(19, dtype=np.uint8),
        'a8': np.arange(23, dtype=np.uint8),
    }

    # On Windows OSError 22 if reusing the same path for memmap ...
    fname = tmpdir.join('test2.mmap').strpath
    numpy_pickle.dump(array_dict, fname, protocol=protocol)
    d_reloaded = numpy_pickle.load(fname, mmap_mode='r')

    for key, memmap in d_reloaded.items():
        assert isinstance(memmap, np.memmap)
        np.testing.assert_array_equal(array_dict[key], memmap)
        assert (
            memmap.ctypes.data % numpy_pickle.NUMPY_ARRAY_ALIGNMENT_BYTES == 0)
        assert memmap.flags.aligned