#-----------------------------------------------------------------------------
#
#  Copyright (c) 2006 by Enthought, Inc.
#  All rights reserved.
#
#  Author: Dave Peterson <dpeterson@enthought.com>
#
#-----------------------------------------------------------------------------

""" A pickle framework that supports unpickling of refactored versions of
    old classes. Part of the AppTools project of the Enthought Tool Suite.

    Beyond refactoring support, there are additional useful capabilities of
    this package:

    - HasTraits objects are fully pickled when sweet_pickle.dump() or
      sweet_pickle.dumps() are called.  This means every trait's value
      is pickled explicitly.  This allows improved self-consistency as
      values within unpickled objects correspond to the default trait
      values of the instance that was pickled rather than the default
      values of the current class definition.
    - Multiple updaters can be registered to update state for any version
      of a class.  This is extremely useful when an instance being
      pickled and unpickled is the result of an addition of one or more
      Trait Categories to a class.  In that case, the Category contributor
      can also register a state function to update the Category's traits
      through versioning.


    We have duplicated the API of the Python 'pickle' module within this
    package. Thus, users can simply import sweet_pickle in places where
    they previously used pickle or cPickle.  For example::

        import cPickle          --->    import apptools.sweet_pickle as pickle
        s = pickle.dumps(obj)           s = pickle.dumps(obj)
        pickle.loads(s)                 pickle.loads(s)

    Pickles generated by this package *can be* unpickled by standard pickle
    or cPickle, though you lose the benefit of the refactoring support during
    unpickling.

    As a result of the above, this package is a drop-in replacement for the
    Python 'pickle' module or 'cPickle' module and should be safe to use even
    if you never version or refactor your own classes.  In fact, we STRONGLY
    RECOMMEND that you use this framework for all of your pickle needs because
    you never know when one of the classes you encounter during unpickling has
    been versioned or otherwise refactored by someone else.

    See module 'pickle' for more basic information about pickling.


    The most common way to benefit from the versioning capabilities of this
    framework is to register class mappings and state modification functions
    with the global updater registry (more detail is below.)  However, you may
    also choose to explicitly instantiate an Unpickler and provide it with your
    own explicit definition of class mappings and state modification functions.
    We do not provide any help on the latter at this time.

    You can register class mappings and state modification functions with the
    global updater registry using the methods provided on the Updater instance
    that is the global registry.  You can get this instance by calling the
    'get_global_registry' function exposed through this package's namespace.

    This framework has been designed so that you can register your class
    mappings and state modification functions during the import of the module
    or package that contained the original class.  However, you may also
    register your mappings and functions at any point such that the they are
    known to the framework prior to, or become known during the attempt to,
    unpickle the class they modify.

    The framework will call a __setstate__ method on the final target class
    of any unpickled instance.  It will not call __setstate__ methods on any
    beginning or intermediate classes within a chain of class mappings.


    A class mapping is used to redirect the unpickling of one class to return
    an instantiation of another class.  The classes can be in different
    modules, and the modules in different packages.  Mappings can be chained.
    For example, given the mappings::

        foo.bar.Bar --> foo.baz.Baz
        foo.baz.Baz --> foo.Foo

    An attempt to unpickle a foo.bar.Bar would actually generate a foo.Foo
    instance.

    A state modification function is called during the execution of the
    __setstate__ method during unpickling of an object of the type and version
    for which it was registered for.  The function must accept a single
    argument which is a state dictionary and must then return the modified
    state dictionary.  Additionally, the function should change the version
    variable within the state to represent the version the new state
    represents.  (The framework will assume an increment of one if this is
    not done.)  The framework ensures that state modification functions
    are chained appropriately to convert through multiple versions and/or class
    mappings.

    Note that refactorings that cause classes to be completed removed from
    the source code can be supported, without breaking unpickling of object
    hierarchies that include an instace of that class, by adding a mapping
    to the Placeholder class in the placeholder module.
"""

##############################################################################
# Implement the Python pickle package API
##############################################################################

# Expose our custom pickler as the standard Unpickler
from .versioned_unpickler import VersionedUnpickler as Unpickler

# Use our custom unpickler to load from files
def load(file, max_pass=-1):
    return Unpickler(file).load(max_pass)

# Use our custom unpickler to load from strings
def loads(str, max_pass=-1):
    from io import BytesIO
    file = BytesIO(str)

    return Unpickler(file).load(max_pass)

# We don't customize the Python pickler, though we do use the cPickle module
# for improved performance.
from pickle import Pickler

# Implement the dump and dumps methods so that all traits in a HasTraits object
# get included in the pickle.
def dump(obj, file, protocol=2):
    _flush_traits(obj)
    from pickle import dump as d
    return d(obj, file, protocol)

def dumps(obj, protocol=2):
    _flush_traits(obj)
    from pickle import dumps as ds
    return ds(obj, protocol)

# We don't customize exceptions so just map to the Python pickle package
from pickle import PickleError, PicklingError, UnpicklingError


##############################################################################
# Allow retrieval of the global registry
##############################################################################

from .global_registry import get_global_registry


##############################################################################
# Expose our Updater class so users can explicitly create their own.
##############################################################################

from .updater import Updater


##############################################################################
# Method to ensure that all traits on a has traits object are included in a
# pickle.
##############################################################################

def _flush_traits(obj):
    if hasattr(obj, 'trait_names'):
        for name, value in list(obj.traits().items()):
            if value.type == 'trait':
                try:
                    getattr(obj, name)
                except AttributeError:
                    pass


### EOF ######################################################################
