basabuuka_prototyp/venv/lib/python3.5/site-packages/hickle/hickle.py

# encoding: utf-8
"""
# hickle.py

Created by Danny Price 2016-02-03.

Hickle is a HDF5 based clone of Pickle. Instead of serializing to a pickle
file, Hickle dumps to a HDF5 file. It is designed to be as similar to pickle in
usage as possible, providing a load() and dump() function.

## Notes

Hickle has two main advantages over Pickle:
1) LARGE PICKLE HANDLING. Unpickling a large pickle is slow, as the Unpickler
reads the entire pickle thing and loads it into memory. In comparison, HDF5
files are designed for large datasets. Things are only loaded when accessed.

2) CROSS PLATFORM SUPPORT. Attempting to unpickle a pickle pickled on Windows
on Linux and vice versa is likely to fail with errors like "Insecure string
pickle". HDF5 files will load fine, as long as both machines have
h5py installed.

"""

from __future__ import absolute_import
import sys
import os
from pkg_resources import get_distribution

import numpy as np
import h5py as h5


from .helpers import get_type_and_data, sort_keys, check_is_iterable, check_iterable_item_type
from .lookup import types_dict, hkl_types_dict, types_not_to_sort, \
    container_types_dict, container_key_types_dict
from .lookup import check_is_ndarray_like


try:
    from exceptions import Exception
    from types import NoneType
except ImportError:
    pass        # above imports will fail in python3

import six
import io

# Import a default 'pickler'
# Not the nicest import code, but should work on Py2/Py3
try:
    import dill as pickle
except ImportError:
    try:
        import cPickle as pickle
    except ImportError:
        import pickle
    except ModuleNotFoundError:
        import pickle
except ModuleNotFoundError:
    import pickle

import warnings

from pkg_resources import get_distribution, DistributionNotFound
try:
    __version__ = get_distribution('hickle').version
except DistributionNotFound:
    __version__ = '0.0.0 - please install via pip/setup.py'

##################
# Error handling #
##################

class FileError(Exception):
    """ An exception raised if the file is fishy """
    def __init__(self):
        return

    def __str__(self):
        return ("Cannot open file. Please pass either a filename "
                "string, a file object, or a h5py.File")


class ClosedFileError(Exception):
    """ An exception raised if the file is fishy """
    def __init__(self):
        return

    def __str__(self):
        return ("HDF5 file has been closed. Please pass either "
                "a filename string, a file object, or an open h5py.File")


class NoMatchError(Exception):
    """ An exception raised if the object type is not understood (or
    supported)"""
    def __init__(self):
        return

    def __str__(self):
        return ("Error: this type of python object cannot be converted into a "
                "hickle.")


class ToDoError(Exception):
    """ An exception raised for non-implemented functionality"""
    def __init__(self):
        return

    def __str__(self):
        return "Error: this functionality hasn't been implemented yet."


class SerializedWarning(UserWarning):
    """ An object type was not understood

    The data will be serialized using pickle.
    """
    pass


######################
# H5PY file wrappers #
######################

class H5GroupWrapper(h5.Group):
    """ Group wrapper that provides a track_times kwarg.

    track_times is a boolean flag that can be set to False, so that two
    files created at different times will have identical MD5 hashes.
    """
    def create_dataset(self, *args, **kwargs):
        kwargs['track_times'] = getattr(self, 'track_times', True)
        return super(H5GroupWrapper, self).create_dataset(*args, **kwargs)

    def create_group(self, *args, **kwargs):
        group = super(H5GroupWrapper, self).create_group(*args, **kwargs)
        group.__class__ = H5GroupWrapper
        group.track_times = getattr(self, 'track_times', True)
        return group


class H5FileWrapper(h5.File):
    """ Wrapper for h5py File that provides a track_times kwarg.

    track_times is a boolean flag that can be set to False, so that two
    files created at different times will have identical MD5 hashes.
    """
    def create_dataset(self, *args, **kwargs):
        kwargs['track_times'] = getattr(self, 'track_times', True)
        return super(H5FileWrapper, self).create_dataset(*args, **kwargs)

    def create_group(self, *args, **kwargs):
        group = super(H5FileWrapper, self).create_group(*args, **kwargs)
        group.__class__ = H5GroupWrapper
        group.track_times = getattr(self, 'track_times', True)
        return group


def file_opener(f, mode='r', track_times=True):
    """ A file opener helper function with some error handling.  This can open
    files through a file object, a h5py file, or just the filename.

    Args:
        f (file, h5py.File, or string): File-identifier, e.g. filename or file object.
        mode (str): File open mode. Only required if opening by filename string.
        track_times (bool): Track time in HDF5; turn off if you want hickling at
                 different times to produce identical files (e.g. for MD5 hash check).

    """

    # Were we handed a file object or just a file name string?
    if six.PY2:
        if isinstance(f, file):
            filename, mode = f.name, f.mode
            f.close()
            h5f = h5.File(filename, mode)
        elif isinstance(f, str) or isinstance(f, unicode):
            filename = f
            h5f = h5.File(filename, mode)
        elif isinstance(f, H5FileWrapper) or isinstance(f, h5._hl.files.File):
            try:
                filename = f.filename
            except ValueError:
                raise ClosedFileError()
            h5f = f
        else:
            print(type(f))
            raise FileError

    else:
        if isinstance(f, io.TextIOWrapper):
            filename, mode = f.name, f.mode
            f.close()
            h5f = h5.File(filename, mode)
        elif isinstance(f, str) or isinstance(f, bytes):
            filename = f
            h5f = h5.File(filename, mode)
        elif isinstance(f, H5FileWrapper) or isinstance(f, h5._hl.files.File):
            try:
                filename = f.filename
            except ValueError:
                raise ClosedFileError()
            h5f = f
        else:
            print(type(f))
            raise FileError


    h5f.__class__ = H5FileWrapper
    h5f.track_times = track_times
    return h5f


###########
# DUMPERS #
###########


def _dump(py_obj, h_group, call_id=0, **kwargs):
    """ Dump a python object to a group within a HDF5 file.

    This function is called recursively by the main dump() function.

    Args:
        py_obj: python object to dump.
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """

    if six.PY2:
        dumpable_dtypes = (bool, int, float, long, complex, str, unicode)
    else:
        dumpable_dtypes = (bool, int, float, complex, bytes, str)

    # Firstly, check if item is a numpy array. If so, just dump it.
    if check_is_ndarray_like(py_obj):
        create_hkl_dataset(py_obj, h_group, call_id, **kwargs)

    # next, check if item is iterable
    elif check_is_iterable(py_obj):
        item_type = check_iterable_item_type(py_obj)

        # item_type == False implies multiple types. Create a dataset
        if item_type is False:
            h_subgroup = create_hkl_group(py_obj, h_group, call_id)
            for ii, py_subobj in enumerate(py_obj):
                _dump(py_subobj, h_subgroup, call_id=ii, **kwargs)

        # otherwise, subitems have same type. Check if subtype is an iterable
        # (e.g. list of lists), or not (e.g. list of ints, which should be treated
        # as a single dataset).
        else:
            if item_type in dumpable_dtypes:
                create_hkl_dataset(py_obj, h_group, call_id, **kwargs)
            else:
                h_subgroup = create_hkl_group(py_obj, h_group, call_id)
                for ii, py_subobj in enumerate(py_obj):
                    _dump(py_subobj, h_subgroup, call_id=ii, **kwargs)

    # item is not iterable, so create a dataset for it
    else:
        create_hkl_dataset(py_obj, h_group, call_id, **kwargs)


def dump(py_obj, file_obj, mode='w', track_times=True, path='/', **kwargs):
    """ Write a pickled representation of obj to the open file object file.

    Args:
    obj (object): python object o store in a Hickle
    file: file object, filename string, or h5py.File object
            file in which to store the object. A h5py.File or a filename is also
            acceptable.
    mode (str): optional argument, 'r' (read only), 'w' (write) or 'a' (append).
            Ignored if file is a file object.
    compression (str): optional argument. Applies compression to dataset. Options: None, gzip,
            lzf (+ szip, if installed)
    track_times (bool): optional argument. If set to False, repeated hickling will produce
            identical files.
    path (str): path within hdf5 file to save data to. Defaults to root /
    """

    try:
        # Open the file
        h5f = file_opener(file_obj, mode, track_times)
        h5f.attrs[b"CLASS"] = b'hickle'
        h5f.attrs[b"VERSION"] = get_distribution('hickle').version
        h5f.attrs[b"type"] = [b'hickle']
        # Log which version of python was used to generate the hickle file
        pv = sys.version_info
        py_ver = "%i.%i.%i" % (pv[0], pv[1], pv[2])
        h5f.attrs[b"PYTHON_VERSION"] = py_ver

        h_root_group = h5f.get(path)

        if h_root_group is None:
            h_root_group = h5f.create_group(path)
            h_root_group.attrs[b"type"] = [b'hickle']

        _dump(py_obj, h_root_group, **kwargs)
        h5f.close()
    except NoMatchError:
        fname = h5f.filename
        h5f.close()
        try:
            os.remove(fname)
        except OSError:
            warnings.warn("Dump failed. Could not remove %s" % fname)
        finally:
            raise NoMatchError


def create_dataset_lookup(py_obj):
    """ What type of object are we trying to pickle?  This is a python
    dictionary based equivalent of a case statement.  It returns the correct
    helper function for a given data type.

    Args:
        py_obj: python object to look-up what function to use to dump to disk

    Returns:
        match: function that should be used to dump data to a new dataset
    """
    t = type(py_obj)
    types_lookup = {dict: create_dict_dataset}
    types_lookup.update(types_dict)

    match = types_lookup.get(t, no_match)

    return match



def create_hkl_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ Create a dataset within the hickle HDF5 file

    Args:
        py_obj: python object to dump.
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.

    """
    #lookup dataset creator type based on python object type
    create_dataset = create_dataset_lookup(py_obj)

    # do the creation
    create_dataset(py_obj, h_group, call_id, **kwargs)


def create_hkl_group(py_obj, h_group, call_id=0):
    """ Create a new group within the hickle file

    Args:
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.

    """
    h_subgroup = h_group.create_group('data_%i' % call_id)
    if six.PY2:
        h_subgroup.attrs["type"] = [str(type(py_obj))]
    else:
        h_subgroup.attrs["type"] = [bytes(str(type(py_obj)), 'ascii')]
    return h_subgroup


def create_dict_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ Creates a data group for each key in dictionary

    Notes:
        This is a very important function which uses the recursive _dump
        method to build up hierarchical data models stored in the HDF5 file.
        As this is critical to functioning, it is kept in the main hickle.py
        file instead of in the loaders/ directory.

    Args:
        py_obj: python object to dump; should be dictionary
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    h_dictgroup = h_group.create_group('data_%i' % call_id)
    h_dictgroup.attrs["type"] = [b'dict']

    for key, py_subobj in py_obj.items():
        if six.PY2:
            if type(key) in (unicode, str):
                h_subgroup = h_dictgroup.create_group(key)
            else:
                h_subgroup = h_dictgroup.create_group(str(key))
        else:
            h_subgroup = h_dictgroup.create_group(str(key))
        h_subgroup.attrs["type"] = [b'dict_item']

        if six.PY2:
            h_subgroup.attrs["key_type"] = [str(type(key))]
        else:
            tk = str(type(key)).encode('utf-8')
            h_subgroup.attrs["key_type"] = [tk]

        _dump(py_subobj, h_subgroup, call_id=0, **kwargs)


def no_match(py_obj, h_group, call_id=0, **kwargs):
    """ If no match is made, raise an exception

    Args:
        py_obj: python object to dump; default if item is not matched.
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    pickled_obj = pickle.dumps(py_obj)
    d = h_group.create_dataset('data_%i' % call_id, data=[pickled_obj])
    d.attrs["type"] = [b'pickle']

    warnings.warn("%s type not understood, data have been serialized" % type(py_obj),
                  SerializedWarning)



#############
## LOADERS ##
#############

class PyContainer(list):
    """ A group-like object into which to load datasets.

    In order to build up a tree-like structure, we need to be able
    to load datasets into a container with an append() method.
    Python tuples and sets do not allow this. This class provides
    a list-like object that be converted into a list, tuple, set or dict.
    """
    def __init__(self):
        super(PyContainer, self).__init__()
        self.container_type = None
        self.name = None
        self.key_type = None

    def convert(self):
        """ Convert from PyContainer to python core data type.

        Returns: self, either as a list, tuple, set or dict
                 (or other type specified in lookup.py)
        """

        if self.container_type in container_types_dict.keys():
            convert_fn = container_types_dict[self.container_type]
            return convert_fn(self)
        if self.container_type == b"dict":
            keys = []
            for item in self:
                key = item.name.split('/')[-1]
                key_type = item.key_type[0]
                if key_type in container_key_types_dict.keys():
                    to_type_fn = container_key_types_dict[key_type]
                    key = to_type_fn(key)
                keys.append(key)

            items = [item[0] for item in self]
            return dict(zip(keys, items))
        else:
            return self

def no_match_load(key):
    """ If no match is made when loading, need to raise an exception
    """
    raise RuntimeError("Cannot load %s data type" % key)
    #pass

def load_dataset_lookup(key):
    """ What type of object are we trying to unpickle?  This is a python
    dictionary based equivalent of a case statement.  It returns the type
    a given 'type' keyword in the hickle file.

    Args:
        py_obj: python object to look-up what function to use to dump to disk

    Returns:
        match: function that should be used to dump data to a new dataset
    """

    match = hkl_types_dict.get(key, no_match_load)

    return match

def load(fileobj, path='/', safe=True):
    """ Load a hickle file and reconstruct a python object

    Args:
        fileobj: file object, h5py.File, or filename string
            safe (bool): Disable automatic depickling of arbitrary python objects.
            DO NOT set this to False unless the file is from a trusted source.
            (see http://www.cs.jhu.edu/~s/musings/pickle.html for an explanation)

        path (str): path within hdf5 file to save data to. Defaults to root /
    """

    try:
        with file_opener(fileobj) as h5f:
            h_root_group = h5f.get(path)
            try:
                assert 'CLASS' in h5f.attrs.keys()
                assert 'VERSION' in h5f.attrs.keys()
                VER = h5f.attrs['VERSION']
                try:
                    VER_MAJOR = int(VER)
                except ValueError:
                    VER_MAJOR = int(VER[0])
                if VER_MAJOR == 1:
                    if six.PY2:
                        warnings.warn("Hickle file versioned as V1, attempting legacy loading...")
                        from . import hickle_legacy
                        return hickle_legacy.load(fileobj, safe)
                    else:
                        raise RuntimeError("Cannot open file. This file was likely"
                                           " created with Python 2 and an old hickle version.")
                elif VER_MAJOR == 2:
                    if six.PY2:
                        warnings.warn("Hickle file appears to be old version (v2), attempting "
                                      "legacy loading...")
                        from . import hickle_legacy2
                        return hickle_legacy2.load(fileobj, safe=safe)
                    else:
                        raise RuntimeError("Cannot open file. This file was likely"
                                           " created with Python 2 and an old hickle version.")
                # There is an unfortunate period of time where hickle 2.1.0 claims VERSION = int(3)
                # For backward compatibility we really need to catch this.
                # Actual hickle v3 files are versioned as A.B.C (e.g. 3.1.0)
                elif VER_MAJOR == 3 and VER == VER_MAJOR:
                    if six.PY2:
                        warnings.warn("Hickle file appears to be old version (v2.1.0), attempting "
                                      "legacy loading...")
                        from . import hickle_legacy2
                        return hickle_legacy2.load(fileobj, safe=safe)
                    else:
                        raise RuntimeError("Cannot open file. This file was likely"
                                           " created with Python 2 and an old hickle version.")
                elif VER_MAJOR >= 3:
                    py_container = PyContainer()
                    py_container.container_type = 'hickle'
                    py_container = _load(py_container, h_root_group)
                    return py_container[0][0]

            except AssertionError:
                if six.PY2:
                    warnings.warn("Hickle file is not versioned, attempting legacy loading...")
                    from . import hickle_legacy
                    return hickle_legacy.load(fileobj, safe)
                else:
                    raise RuntimeError("Cannot open file. This file was likely"
                                       " created with Python 2 and an old hickle version.")
    finally:
        if 'h5f' in locals():
            # Check if file is open, and if so, close it.
            if h5f.fid.valid:
                h5f.close()

def load_dataset(h_node):
    """ Load a dataset, converting into its correct python type

    Args:
        h_node (h5py dataset): h5py dataset object to read

    Returns:
        data: reconstructed python object from loaded data
    """
    py_type, data = get_type_and_data(h_node)

    try:
        load_fn = load_dataset_lookup(py_type)
        return load_fn(h_node)
    except:
        raise
        #raise RuntimeError("Hickle type %s not understood." % py_type)

def _load(py_container, h_group):
    """ Load a hickle file

    Recursive funnction to load hdf5 data into a PyContainer()

    Args:
        py_container (PyContainer): Python container to load data into
        h_group (h5 group or dataset): h5py object, group or dataset, to spider
                and load all datasets.
    """

    group_dtype   = h5._hl.group.Group
    dataset_dtype = h5._hl.dataset.Dataset

    #either a file, group, or dataset
    if isinstance(h_group, H5FileWrapper) or isinstance(h_group, group_dtype):

        py_subcontainer = PyContainer()
        try:
            py_subcontainer.container_type = bytes(h_group.attrs['type'][0])
        except KeyError:
            raise
            #py_subcontainer.container_type = ''
        py_subcontainer.name = h_group.name

        if py_subcontainer.container_type == b'dict_item':
            py_subcontainer.key_type = h_group.attrs['key_type']

        if py_subcontainer.container_type not in types_not_to_sort:
            h_keys = sort_keys(h_group.keys())
        else:
            h_keys = h_group.keys()

        for h_name in h_keys:
            h_node = h_group[h_name]
            py_subcontainer = _load(py_subcontainer, h_node)

        sub_data = py_subcontainer.convert()
        py_container.append(sub_data)

    else:
        # must be a dataset
        subdata = load_dataset(h_group)
        py_container.append(subdata)

    return py_container