<commit_before>"""
A pickler to save numpy arrays in separate .npy files.
"""

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

import pickle
import traceback
import sys
import os
import shutil
import tempfile
import zipfile
import warnings

if sys.version_info[0] == 3:
    from pickle import _Unpickler as Unpickler
    from cStringIO import StringIO as BytesIO
else:
    from io import BytesIO
    from pickle import Unpickler


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


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

        The only thing this object does, is store the filename in wich
        the array has been persisted.
    """
    def __init__(self, filename, subclass=None):
        self.filename = filename
        self.subclass = subclass


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

class NumpyPickler(pickle.Pickler):
    """ A pickler subclass that extracts ndarrays and saves them in .npy
        files outside of the pickle.
    """

    def __init__(self, filename):
        self._filename = filename
        self._filenames = [filename, ]
        self.file = open(filename, 'wb')
        # Count the number of npy files that we have created:
        self._npy_counter = 0
        pickle.Pickler.__init__(self, self.file,
                                protocol=pickle.HIGHEST_PROTOCOL)
        # delayed import of numpy, to avoid tight coupling
        try:
            import numpy as np
        except ImportError:
            np = None
        self.np = np

    def save(self, obj):
        """ Subclass the save method, to save ndarray subclasses in npy
            files, rather than pickling them. Of course, this is a
            total abuse of the Pickler class.
        """
        if self.np is not None and type(obj) in (self.np.ndarray,
                                                 self.np.matrix, self.np.memmap):
            self._npy_counter += 1
            try:
                filename = '%s_%02i.npy' % (self._filename,
                                            self._npy_counter)
                self.np.save(filename, obj)
                self._filenames.append(filename)
                obj = NDArrayWrapper(os.path.basename(filename),
                                     type(obj))
            except:
                self._npy_counter -= 1
                # XXX: We should have a logging mechanism
                print 'Failed to save %s to .npy file:\n%s' % (
                        type(obj),
                        traceback.format_exc())
        pickle.Pickler.save(self, obj)


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

    def __init__(self, filename, file_handle=None, mmap_mode=None):
        self._filename = os.path.basename(filename)
        self.mmap_mode = mmap_mode
        self._dirname = os.path.dirname(filename)
        if file_handle is None:
            file_handle = self._open_file(self._filename)
            if isinstance(file_handle, basestring):
                # To handle memmap, we need to have file names
                file_handle = open(file_handle, 'rb')
        self.file_handle = file_handle
        Unpickler.__init__(self, self.file_handle)
        try:
            import numpy as np
        except ImportError:
            np = None
        self.np = np

    def _open_file(self, name):
        "Return the path of the given file in our store"
        return os.path.join(self._dirname, name)

    def load_build(self):
        """ This method is called to set the state of a newly created
            object.

            We capture it to replace our place-holder objects,
            NDArrayWrapper, by the array we are interested in. We
            replace directly in the stack of pickler.
        """
        Unpickler.load_build(self)
        if isinstance(self.stack[-1], NDArrayWrapper):
            if self.np is None:
                raise ImportError('Trying to unpickle an ndarray, '
                        "but numpy didn't import correctly")
            nd_array_wrapper = self.stack.pop()
            if self.np.__version__ >= '1.3':
                array = self.np.load(
                                self._open_file(nd_array_wrapper.filename),
                                mmap_mode=self.mmap_mode)
            else:
                # Numpy does not have mmap_mode before 1.3
                array = self.np.load(
                                self._open_file(nd_array_wrapper.filename),
                                mmap_mode=self.mmap_mode)
            if not nd_array_wrapper.subclass is self.np.ndarray:
                # We need to reconstruct another subclass
                new_array = self.np.core.multiarray._reconstruct(
                        nd_array_wrapper.subclass, (0,), 'b')
                new_array.__array_prepare__(array)
                array = new_array
            self.stack.append(array)

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


class ZipNumpyUnpickler(NumpyUnpickler):
    """ A subclass of our Unpickler to unpickle on the fly from zips.
    """

    def __init__(self, file_handle):
        kwargs = dict(compression=zipfile.ZIP_DEFLATED)
        if sys.version_info >= (2, 5):
            kwargs['allowZip64'] = True
        self._zip_file = zipfile.ZipFile(file_handle, **kwargs)
        NumpyUnpickler.__init__(self, 'joblib_dump.pkl',
                                mmap_mode=None)

    def _open_file(self, name):
        "Return the path of the given file in our store"
        decompression_buffer = BytesIO(
                self._zip_file.read(os.path.join('dump_file', name)))
        return decompression_buffer

<commit_msg>ENH: Improved memory usage in ZipUnpickler<commit_after>"""
A pickler to save numpy arrays in separate .npy files.
"""

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

import pickle
import traceback
import sys
import os
import shutil
import tempfile
import zipfile
import warnings

if sys.version_info[0] == 3:
    from pickle import _Unpickler as Unpickler
    from cStringIO import StringIO as BytesIO
else:
    from io import BytesIO
    from pickle import Unpickler


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


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

        The only thing this object does, is store the filename in wich
        the array has been persisted.
    """
    def __init__(self, filename, subclass=None):
        self.filename = filename
        self.subclass = subclass


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

class NumpyPickler(pickle.Pickler):
    """ A pickler subclass that extracts ndarrays and saves them in .npy
        files outside of the pickle.
    """

    def __init__(self, filename):
        self._filename = filename
        self._filenames = [filename, ]
        self.file = open(filename, 'wb')
        # Count the number of npy files that we have created:
        self._npy_counter = 0
        pickle.Pickler.__init__(self, self.file,
                                protocol=pickle.HIGHEST_PROTOCOL)
        # delayed import of numpy, to avoid tight coupling
        try:
            import numpy as np
        except ImportError:
            np = None
        self.np = np

    def save(self, obj):
        """ Subclass the save method, to save ndarray subclasses in npy
            files, rather than pickling them. Of course, this is a
            total abuse of the Pickler class.
        """
        if self.np is not None and type(obj) in (self.np.ndarray,
                                                 self.np.matrix, self.np.memmap):
            self._npy_counter += 1
            try:
                filename = '%s_%02i.npy' % (self._filename,
                                            self._npy_counter)
                self.np.save(filename, obj)
                self._filenames.append(filename)
                obj = NDArrayWrapper(os.path.basename(filename),
                                     type(obj))
            except:
                self._npy_counter -= 1
                # XXX: We should have a logging mechanism
                print 'Failed to save %s to .npy file:\n%s' % (
                        type(obj),
                        traceback.format_exc())
        pickle.Pickler.save(self, obj)


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

    def __init__(self, filename, file_handle=None, mmap_mode=None):
        self._filename = os.path.basename(filename)
        self.mmap_mode = mmap_mode
        self._dirname = os.path.dirname(filename)
        if file_handle is None:
            file_handle = self._open_file(self._filename)
            if isinstance(file_handle, basestring):
                # To handle memmap, we need to have file names
                file_handle = open(file_handle, 'rb')
        self.file_handle = file_handle
        Unpickler.__init__(self, self.file_handle)
        try:
            import numpy as np
        except ImportError:
            np = None
        self.np = np

    def _open_file(self, name):
        "Return the path of the given file in our store"
        return os.path.join(self._dirname, name)

    def load_build(self):
        """ This method is called to set the state of a newly created
            object.

            We capture it to replace our place-holder objects,
            NDArrayWrapper, by the array we are interested in. We
            replace directly in the stack of pickler.
        """
        Unpickler.load_build(self)
        if isinstance(self.stack[-1], NDArrayWrapper):
            if self.np is None:
                raise ImportError('Trying to unpickle an ndarray, '
                        "but numpy didn't import correctly")
            nd_array_wrapper = self.stack.pop()
            obj = self._open_file(nd_array_wrapper.filename)
            if not isinstance(obj, basestring):
                # Must be a file-like object
                # read_array can handle them
                array = self.np.lib.format.read_array(obj)
            elif self.np.__version__ >= '1.3':
                array = self.np.load(obj,
                                mmap_mode=self.mmap_mode)
            else:
                # Numpy does not have mmap_mode before 1.3
                array = self.np.load(obj)
            if (not nd_array_wrapper.subclass is self.np.ndarray 
                    and not nd_array_wrapper.subclass is self.np.memmap):
                # We need to reconstruct another subclass
                new_array = self.np.core.multiarray._reconstruct(
                        nd_array_wrapper.subclass, (0,), 'b')
                new_array.__array_prepare__(array)
                array = new_array
            self.stack.append(array)

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


class ZipNumpyUnpickler(NumpyUnpickler):
    """ A subclass of our Unpickler to unpickle on the fly from zips.
    """

    def __init__(self, file_handle):
        kwargs = dict(compression=zipfile.ZIP_DEFLATED)
        if sys.version_info >= (2, 5):
            kwargs['allowZip64'] = True
        self._zip_file = zipfile.ZipFile(file_handle, **kwargs)
        NumpyUnpickler.__init__(self, 'joblib_dump.pkl',
                                mmap_mode=None)

    def _open_file(self, name):
        "Return the path of the given file in our store"
        if name.endswith('.npy'):
            # Only numpy's reader can handle file-like objects
            return self._zip_file.open(os.path.join('dump_file', name),
                                       'r')
        else:
            decompression_buffer = \
                    self._zip_file.read(os.path.join('dump_file', name))
            decompression_buffer = BytesIO(decompression_buffer)
        return decompression_buffer


{'fraction_non_alphanumeric': 0.1020096463022508, 'fraction_numerical': 0.003536977491961415, 'mean_word_length': 1.885204081632653, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 11, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '476227', 'n_tokens_mistral': 3463, 'n_tokens_neox': 3284, 'n_words': 1005}