"""This module provides code that allows one to pickle the state of a

Python object to a dictionary.



The motivation for this is simple.  The standard Python

pickler/unpickler is best used to pickle simple objects and does not

work too well for complex code.  Specifically, there are two major

problems (1) the pickle file format is not easy to edit with a text

editor and (2) when a pickle is unpickled, it creates all the

necessary objects and sets the state of these objects.



Issue (2) might not appear to be a problem.  However, often, the

determination of the entire 'state' of an application requires the

knowledge of the state of many objects that are not really in the

users concern.  The user would ideally like to pickle just what he

thinks is relevant.  Now, given that the user is not going to save the

entire state of the application, the use of pickle is insufficient

since the state is no longer completely known (or worth knowing).  The

default `Unpickler` recreates the objects and the typical

implementation of `__setstate__` is usually to simply update the

object's `__dict__` attribute.  This is inadequate because the pickled

information is taken out of the real context when it was saved.



The `StatePickler` basically pickles the 'state' of an object into a

large dictionary.  This pickled data may be easily unpickled and

modified on the interpreter or edited with a text editor

(`pprint.saferepr` is a friend).  The second problem is also

eliminated.  When this state is unpickled using `StateUnpickler`, what

you get is a special dictionary (a `State` instance).  This allows one

to navigate the state just like the original object.  Its up to the

user to create any new objects and set their states using this

information.  This allows for a lot of flexibility while allowing one

to save and set the state of (almost) any Python object.



The `StateSetter` class helps set the state of a known instance.  When

setting the state of an instance it checks to see if there is a

`__set_pure_state__` method that in turn calls `StateSetter.set`

appropriately.



Additionally, there is support for versioning.  The class' version is

obtain from the `__version__` class attribute.  This version along

with the versions of the bases of a class is embedded into the

metadata of the state and stored.  By using `version_registry.py` a

user may register a handler for a particular class and module.  When

the state of an object is set using `StateSetter.set_state`, then

these handlers are called in reverse order of their MRO.  This gives

the handler an opportunity to upgrade the state depending on its

version.  Builtin classes are not scanned for versions.  If a class

has no version, then by default it is assumed to be -1.





Example::



  >>> class A:

  ...    def __init__(self):

  ...        self.a = 'a'

  ...

  >>> a = A()

  >>> a.a = 100

  >>> from . import state_pickler

  >>> s = state_pickler.dumps(a)               # Dump the state of `a`.

  >>> state = state_pickler.loads_state(s)     # Get the state back.

  >>> b = state_pickler.create_instance(state) # Create the object.

  >>> state_pickler.set_state(b, state)        # Set the object's state.

  >>> assert b.a == 100





Features

--------



 - The output is a plain old dictionary so is easy to parse, edit etc.

 - Handles references to avoid duplication.

 - Gzips Numeric arrays when dumping them.

 - Support for versioning.





Caveats

-------



 - Does not pickle a whole bunch of stuff including code objects and

   functions.

 - The output is a pure dictionary and does not contain instances.  So

   using this *as it is* in `__setstate__` will not work.  Instead

   define a `__set_pure_state__` and use the `StateSetter` class or

   the `set_state` function provided by this module.





Notes

-----



  Browsing the code from XMarshaL_ and pickle.py proved useful for

  ideas.  None of the code is taken from there though.



.. _XMarshaL:  http://www.dezentral.de/soft/XMarshaL



"""

# Author: Prabhu Ramachandran <prabhu_r@users.sf.net>

# Copyright (c) 2005-2015, Enthought, Inc.

# License: BSD Style.



# Standard library imports.

import base64

import sys

import types

import pickle

import gzip

from io import BytesIO, StringIO



import numpy



# Local imports.

from . import version_registry

from .file_path import FilePath



PY_VER = sys.version_info[0]

NumpyArrayType = type(numpy.array([]))





def gzip_string(data):

    """Given a string (`data`) this gzips the string and returns it.

    """

    s = BytesIO()

    writer = gzip.GzipFile(mode='wb', fileobj=s)

    writer.write(data)

    writer.close()

    s.seek(0)

    return s.read()





def gunzip_string(data):

    """Given a gzipped string (`data`) this unzips the string and

    returns it.

    """

    if PY_VER== 2 or (bytes is not str and type(data) is bytes):

        s = BytesIO(data)

    else:

        s = StringIO(data)

    writer = gzip.GzipFile(mode='rb', fileobj=s)

    data = writer.read()

    writer.close()

    return data



class StatePicklerError(Exception):

    pass



class StateUnpicklerError(Exception):

    pass



class StateSetterError(Exception):

    pass



######################################################################

# `State` class

######################################################################

class State(dict):

    """Used to encapsulate the state of an instance in a very

    convenient form.  The '__metadata__' attribute/key is a dictionary

    that has class specific details like the class name, module name

    etc.

    """

    def __init__(self, **kw):

        dict.__init__(self, **kw)

        self.__dict__ = self



######################################################################

# `StateDict` class

######################################################################

class StateDict(dict):

    """Used to encapsulate a dictionary stored in a `State` instance.

    The has_instance attribute specifies if the dict has an instance

    embedded in it.

    """

    def __init__(self, **kw):

        dict.__init__(self, **kw)

        self.has_instance = False



######################################################################

# `StateList` class

######################################################################

class StateList(list):

    """Used to encapsulate a list stored in a `State` instance.  The

    has_instance attribute specifies if the list has an instance

    embedded in it.

    """

    def __init__(self, seq=None):

        if seq:

            list.__init__(self, seq)

        else:

            list.__init__(self)

        self.has_instance = False



######################################################################

# `StateTuple` class

######################################################################

class StateTuple(tuple):

    """Used to encapsulate a tuple stored in a `State` instance.  The

    has_instance attribute specifies if the tuple has an instance

    embedded in it.

    """

    def __new__(cls, seq=None):

        if seq:

            obj = super(StateTuple, cls).__new__(cls, tuple(seq))

        else:

            obj = super(StateTuple, cls).__new__(cls)

        obj.has_instance = False

        return obj





######################################################################

# `StatePickler` class

######################################################################

class StatePickler:

    """Pickles the state of an object into a dictionary.  The

    dictionary is itself either saved as a pickled file (`dump`) or

    pickled string (`dumps`).  Alternatively, the `dump_state` method

    will return the dictionary that is pickled.



    The format of the state dict is quite strightfoward.  Basic types

    (bool, int, long, float, complex, None, string and unicode) are

    represented as they are.  Everything else is stored as a

    dictionary containing metadata information on the object's type

    etc. and also the actual object in the 'data' key.  For example::



        >>> p = StatePickler()

        >>> p.dump_state(1)

        1

        >>> l = [1,2.0, None, [1,2,3]]

        >>> p.dump_state(l)

        {'data': [1, 2.0, None, {'data': [1, 2, 3], 'type': 'list', 'id': 1}],

         'id': 0,

         'type': 'list'}



    Classes are also represented similarly.  The state in this case is

    obtained from the `__getstate__` method or from the `__dict__`.

    Here is an example::



        >>> class A:

        ...     __version__ = 1  # State version

        ...     def __init__(self):

        ...         self.attribute = 1

        ...

        >>> a = A()

        >>> p = StatePickler()

        >>> p.dump_state(a)

        {'class_name': 'A',

         'data': {'data': {'attribute': 1}, 'type': 'dict', 'id': 2},

         'id': 0,

         'initargs': {'data': (), 'type': 'tuple', 'id': 1},

         'module': '__main__',

         'type': 'instance',

         'version': [(('A', '__main__'), 1)]}



    When pickling data, references are taken care of.  Numeric arrays

    can be pickled and are stored as a gzipped base64 encoded string.



    """

    def __init__(self):

        self._clear()

        type_map = {bool: self._do_basic_type,

                    complex: self._do_basic_type,

                    float: self._do_basic_type,

                    int: self._do_basic_type,

                    type(None): self._do_basic_type,

                    str: self._do_basic_type,

                    bytes: self._do_basic_type,

                    tuple: self._do_tuple,

                    list: self._do_list,

                    dict: self._do_dict,

                    NumpyArrayType: self._do_numeric,

                    State: self._do_state,

                    }

        if PY_VER == 2:

            type_map[int] = self._do_basic_type

            type_map[str] = self._do_basic_type

        self.type_map = type_map



    def dump(self, value, file):

        """Pickles the state of the object (`value`) into the passed

        file.

        """

        try:

            # Store the file name we are writing to so we can munge

            # file paths suitably.

            self.file_name = file.name

        except AttributeError:

            pass

        pickle.dump(self._do(value), file)



    def dumps(self, value):

        """Pickles the state of the object (`value`) and returns a

        string.

        """

        return pickle.dumps(self._do(value))



    def dump_state(self, value):

        """Returns a dictionary or a basic type representing the

        complete state of the object (`value`).



        This value is pickled by the `dump` and `dumps` methods.

        """

        return self._do(value)



    ######################################################################

    # Non-public methods

    ######################################################################

    def _clear(self):

        # Stores the file name of the file being used to dump the

        # state.  This is used to change any embedded paths relative

        # to the saved file.

        self.file_name = ''

        # Caches id's to handle references.

        self.obj_cache = {}

        # Misc cache to cache things that are not persistent.  For

        # example, object.__getstate__()/__getinitargs__() usually

        # returns a copy of a dict/tuple that could possibly be reused

        # on another object's __getstate__.  Caching these prevents

        # some wierd problems with the `id` of the object.

        self._misc_cache = []



    def _flush_traits(self, obj):

        """Checks if the object has traits and ensures that the traits

        are set in the `__dict__` so we can pickle it.

        """

        # Not needed with Traits3.

        return



    def _do(self, obj):

        obj_type = type(obj)

        key = self._get_id(obj)

        if key in self.obj_cache:

            return self._do_reference(obj)

        elif obj_type in self.type_map:

            return self.type_map[obj_type](obj)

        elif isinstance(obj, tuple):

            # Takes care of StateTuples.

            return self._do_tuple(obj)

        elif isinstance(obj, list):

            # Takes care of TraitListObjects.

            return self._do_list(obj)

        elif isinstance(obj, dict):

            # Takes care of TraitDictObjects.

            return self._do_dict(obj)

        elif hasattr(obj, '__dict__'):

            return self._do_instance(obj)



    def _get_id(self, value):

        try:

            key = hash(value)

        except TypeError:

            key = id(value)

        return key



    def _register(self, value):

        key = self._get_id(value)

        cache = self.obj_cache

        idx = len(cache)

        cache[key] = idx

        return idx



    def _do_basic_type(self, value):

        return value



    def _do_reference(self, value):

        key = self._get_id(value)

        idx = self.obj_cache[key]

        return dict(type='reference', id=idx, data=None)



    def _do_instance(self, value):

        # Flush out the traits.

        self._flush_traits(value)



        # Setup the relative paths of FilePaths before dumping.

        if self.file_name and isinstance(value, FilePath):

            value.set_relative(self.file_name)



        # Get the initargs.

        args = ()

        if hasattr(value, '__getinitargs__') and value.__getinitargs__:

            args = value.__getinitargs__()



        # Get the object state.

        if hasattr(value, '__get_pure_state__'):

            state = value.__get_pure_state__()

        elif hasattr(value, '__getstate__'):

            state = value.__getstate__()

        else:

            state = value.__dict__



        state.pop('__traits_version__', None)



        # Cache the args and state since they are likely to be gc'd.

        self._misc_cache.extend([args, state])

        # Register and process.

        idx = self._register(value)

        args_data = self._do(args)

        data = self._do(state)



        # Get the version of the object.

        version = version_registry.get_version(value)

        module = value.__class__.__module__

        class_name = value.__class__.__name__



        return dict(type='instance',

                    module=module,

                    class_name=class_name,

                    version=version,

                    id=idx,

                    initargs=args_data,

                    data=data)



    def _do_state(self, value):

        metadata = value.__metadata__

        args = metadata.get('initargs')

        state = dict(value)

        state.pop('__metadata__')



        self._misc_cache.extend([args, state])



        idx = self._register(value)

        args_data = self._do(args)

        data = self._do(state)



        return dict(type='instance',

                    module=metadata['module'],

                    class_name=metadata['class_name'],

                    version=metadata['version'],

                    id=idx,

                    initargs=args_data,

                    data=data)



    def _do_tuple(self, value):

        idx = self._register(value)

        data = tuple([self._do(x) for x in value])

        return dict(type='tuple', id=idx, data=data)



    def _do_list(self, value):

        idx = self._register(value)

        data = [self._do(x) for x in value]

        return dict(type='list', id=idx, data=data)



    def _do_dict(self, value):

        idx = self._register(value)

        vals = [self._do(x) for x in list(value.values())]

        data = dict(list(zip(list(value.keys()), vals)))

        return dict(type='dict', id=idx, data=data)



    def _do_numeric(self, value):

        idx = self._register(value)

        if PY_VER > 2:

            data = base64.encodebytes(gzip_string(numpy.ndarray.dumps(value)))

        else:

            data = base64.encodestring(gzip_string(numpy.ndarray.dumps(value)))

        return dict(type='numeric', id=idx, data=data)







######################################################################

# `StateUnpickler` class

######################################################################

class StateUnpickler:

    """Unpickles the state of an object saved using StatePickler.



    Please note that unlike the standard Unpickler, no instances of

    any user class are created.  The data for the state is obtained

    from the file or string, reference objects are setup to refer to

    the same state value and this state is returned in the form

    usually in the form of a dictionary.  For example::



        >>> class A:

        ...     def __init__(self):

        ...         self.attribute = 1

        ...

        >>> a = A()

        >>> p = StatePickler()

        >>> s = p.dumps(a)

        >>> up = StateUnpickler()

        >>> state = up.loads_state(s)

        >>> state.__class__.__name__

        'State'

        >>> state.attribute

        1

        >>> state.__metadata__

        {'class_name': 'A',

         'has_instance': True,

         'id': 0,

         'initargs': (),

         'module': '__main__',

         'type': 'instance',

         'version': [(('A', '__main__'), -1)]}



    Note that the state is actually a `State` instance and is

    navigable just like the original object.  The details of the

    instance are stored in the `__metadata__` attribute.  This is

    highly convenient since it is possible for someone to view and

    modify the state very easily.

    """



    def __init__(self):

        self._clear()

        self.type_map = {'reference': self._do_reference,

                         'instance': self._do_instance,

                         'tuple': self._do_tuple,

                         'list': self._do_list,

                         'dict': self._do_dict,

                         'numeric': self._do_numeric,

                         }



    def load_state(self, file):

        """Returns the state of an object loaded from the pickled data

        in the given file.

        """

        try:

            self.file_name = file.name

        except AttributeError:

            pass

        data = pickle.load(file)

        result = self._process(data)

        return result



    def loads_state(self, string):

        """Returns the state of an object loaded from the pickled data

        in the given string.

        """

        data = pickle.loads(string)

        result = self._process(data)

        return result



    ######################################################################

    # Non-public methods

    ######################################################################

    def _clear(self):

        # The file from which we are being loaded.

        self.file_name = ''

        # Cache of the objects.

        self._obj_cache = {}

        # Paths to the instances.

        self._instances = []

        # Caches the references.

        self._refs = {}

        # Numeric arrays.

        self._numeric = {}



    def _set_has_instance(self, obj, value):

        if isinstance(obj, State):

            obj.__metadata__['has_instance'] = value

        elif isinstance(obj, (StateDict, StateList, StateTuple)):

            obj.has_instance = value



    def _process(self, data):

        result = self._do(data)



        # Setup all the Numeric arrays.  Do this first since

        # references use this.

        for key, (path, val) in list(self._numeric.items()):

            if isinstance(result, StateTuple):

                result = list(result)

                exec('result%s = val'%path)

                result = StateTuple(result)

            else:

                exec('result%s = val'%path)



        # Setup the references so they really are references.

        for key, paths in list(self._refs.items()):

            for path in paths:

                x = self._obj_cache[key]

                if isinstance(result, StateTuple):

                    result = list(result)

                    exec('result%s = x'%path)

                    result = StateTuple(result)

                else:

                    exec('result%s = x'%path)

                # if the reference is to an instance append its path.

                if isinstance(x, State):

                    self._instances.append(path)



        # Now setup the 'has_instance' attribute.  If 'has_instance'

        # is True then the object contains an instance somewhere

        # inside it.

        for path in self._instances:

            pth = path

            while pth:

                ns = {'result': result}

                exec('val = result%s'%pth, ns, ns)

                self._set_has_instance(ns['val'], True)

                end = pth.rfind('[')

                pth = pth[:end]

            # Now make sure that the first element also has_instance.

            self._set_has_instance(result, True)

        return result



    def _do(self, data, path=''):

        if type(data) is dict:

            return self.type_map[data['type']](data, path)

        else:

            return data



    def _do_reference(self, value, path):

        id = value['id']

        if id in self._refs:

            self._refs[id].append(path)

        else:

            self._refs[id] = [path]

        return State(__metadata__=value)



    def _handle_file_path(self, value):

        if (value['class_name'] == 'FilePath') and \

           ('file_path' in value['module']) and \

           self.file_name:

            data = value['data']['data']

            fp = FilePath(data['rel_pth'])

            fp.set_absolute(self.file_name)

            data['abs_pth'] = fp.abs_pth



    def _do_instance(self, value, path):

        self._instances.append(path)

        initargs = self._do(value['initargs'],

                            path + '.__metadata__["initargs"]')

        # Handle FilePaths.

        self._handle_file_path(value)



        d = self._do(value['data'], path)

        md = dict(type='instance',

                  module=value['module'],

                  class_name=value['class_name'],

                  version=value['version'],

                  id=value['id'],

                  initargs=initargs,

                  has_instance=True)

        result = State(**d)

        result.__metadata__ = md

        self._obj_cache[value['id']] = result

        return result



    def _do_tuple(self, value, path):

        res = []

        for i, x in enumerate(value['data']):

            res.append(self._do(x, path + '[%d]'%i))

        result = StateTuple(res)

        self._obj_cache[value['id']] = result

        return result



    def _do_list(self, value, path):

        result = StateList()

        for i, x in enumerate(value['data']):

            result.append(self._do(x, path + '[%d]'%i))

        self._obj_cache[value['id']] = result

        return result



    def _do_dict(self, value, path):

        result = StateDict()

        for key, val in list(value['data'].items()):

            result[key] = self._do(val, path + '["%s"]'%key)

        self._obj_cache[value['id']] = result

        return result



    def _do_numeric(self, value, path):

        if PY_VER > 2:

            data = value['data']

            if isinstance(data, str):

                data = value['data'].encode('utf-8')

            junk = gunzip_string(base64.decodebytes(data))

            result = numpy.loads(junk, encoding='bytes')

        else:

            junk = gunzip_string(value['data'].decode('base64'))

            result = numpy.loads(junk)

        self._numeric[value['id']] = (path, result)

        self._obj_cache[value['id']] = result

        return result





######################################################################

# `StateSetter` class

######################################################################

class StateSetter:

    """This is a convenience class that helps a user set the

    attributes of an object given its saved state.  For instances it

    checks to see if a `__set_pure_state__` method exists and calls

    that when it sets the state.

    """

    def __init__(self):

        # Stores the ids of instances already done.

        self._instance_ids = []

        self.type_map = {State: self._do_instance,

                         StateTuple: self._do_tuple,

                         StateList: self._do_list,

                         StateDict: self._do_dict,

                         }



    def set(self, obj, state, ignore=None, first=None, last=None):

        """Sets the state of the object.



        This is to be used as a means to simplify loading the state of

        an object from its `__setstate__` method using the dictionary

        describing its state.  Note that before the state is set, the

        registered handlers for the particular class are called in

        order to upgrade the version of the state to the latest

        version.



        Parameters

        ----------



        - obj : `object`



          The object whose state is to be set.  If this is `None`

          (default) then the object is created.



        - state : `dict`



          The dictionary representing the state of the object.



        - ignore : `list(str)`



          The list of attributes specified in this list are ignored

          and the state of these attributes are not set (this excludes

          the ones specified in `first` and `last`).  If one specifies

          a '*' then all attributes are ignored except the ones

          specified in `first` and `last`.



        - first : `list(str)`



          The list of attributes specified in this list are set first (in

          order), before any other attributes are set.



        - last : `list(str)`



          The list of attributes specified in this list are set last (in

          order), after all other attributes are set.



        """

        if (not isinstance(state, State)) and \

               state.__metadata__['type'] != 'instance':

            raise StateSetterError(

                'Can only set the attributes of an instance.'

            )



        # Upgrade the state to the latest using the registry.

        self._update_and_check_state(obj, state)



        self._register(obj)



        # This wierdness is needed since the state's own `keys` might

        # be set to something else.

        state_keys = list(dict.keys(state))

        state_keys.remove('__metadata__')



        if first is None:

            first = []

        if last is None:

            last = []



        # Remove all the ignored keys.

        if ignore:

            if '*' in ignore:

                state_keys = first + last

            else:

                for name in ignore:

                    try:

                        state_keys.remove(name)

                    except KeyError:

                        pass



        # Do the `first` attributes.

        for key in first:

            state_keys.remove(key)

            self._do(obj, key, state[key])



        # Remove the `last` attributes.

        for key in last:

            state_keys.remove(key)



        # Set the remaining attributes.

        for key in state_keys:

            self._do(obj, key, state[key])



        # Do the last ones in order.

        for key in last:

            self._do(obj, key, state[key])



    ######################################################################

    # Non-public methods.

    ######################################################################

    def _register(self, obj):

        idx = id(obj)

        if idx not in self._instance_ids:

            self._instance_ids.append(idx)



    def _is_registered(self, obj):

        return (id(obj) in self._instance_ids)



    def _has_instance(self, value):

        """Given something (`value`) that is part of the state this

        returns if the value has an instance embedded in it or not.

        """

        if isinstance(value, State):

            return True

        elif isinstance(value, (StateDict, StateList, StateTuple)):

            return value.has_instance

        return False



    def _get_pure(self, value):

        """Returns the Python representation of the object (usually a

        list, tuple or dict) that has no instances embedded within it.

        """

        result = value

        if self._has_instance(value):

            raise StateSetterError(

                'Value has an instance: %s'%value

            )

        if isinstance(value, (StateList, StateTuple)):

            result = [self._get_pure(x) for x in value]

            if isinstance(value, StateTuple):

                result = tuple(result)

        elif isinstance(value, StateDict):

            result = {}

            for k, v in list(value.items()):

                result[k] = self._get_pure(v)

        return result



    def _update_and_check_state(self, obj, state):

        """Updates the state from the registry and then checks if the

        object and state have same class.

        """

        # Upgrade this state object to the latest using the registry.

        # This is done before testing because updating may change the

        # class name/module.

        version_registry.registry.update(state)



        # Make sure object and state have the same class and module names.

        metadata = state.__metadata__

        cls = obj.__class__

        if (metadata['class_name'] != cls.__name__):

            raise StateSetterError(

                'Instance (%s) and state (%s) do not have the same class'\

                ' name!'%(cls.__name__, metadata['class_name'])

            )

        if (metadata['module'] != cls.__module__):

            raise StateSetterError(

                'Instance (%s) and state (%s) do not have the same module'\

                ' name!'%(cls.__module__, metadata['module'])

            )



    def _do(self, obj, key, value):

        try:

            attr = getattr(obj, key)

        except AttributeError:

            raise StateSetterError(

                'Object %s does not have an attribute called: %s'%(obj, key)

            )



        if isinstance(value, (State, StateDict, StateList, StateTuple)):

            # Special handlers are needed.

            if not self._has_instance(value):

                result = self._get_pure(value)

                setattr(obj, key, result)

            elif isinstance(value, StateTuple):

                setattr(obj, key, self._do_tuple(getattr(obj, key), value))

            else:

                self._do_object(getattr(obj, key), value)

        else:

            setattr(obj, key, value)



    def _do_object(self, obj, state):

        self.type_map[state.__class__](obj, state)



    def _do_instance(self, obj, state):

        if self._is_registered(obj):

            return

        else:

            self._register(obj)



        metadata = state.__metadata__

        if hasattr(obj, '__set_pure_state__'):

            self._update_and_check_state(obj, state)

            obj.__set_pure_state__(state)

        elif 'tvtk_classes' in metadata['module']:

            self._update_and_check_state(obj, state)

            tmp = self._get_pure(StateDict(**state))

            del tmp['__metadata__']

            obj.__setstate__(tmp)

        else:

            # No need to update or check since `set` does it for us.

            self.set(obj, state)



    def _do_tuple(self, obj, state):

        if not self._has_instance(state):

            return self._get_pure(state)

        else:

            result = list(obj)

            self._do_list(result, state)

            return tuple(result)



    def _do_list(self, obj, state):

        if len(obj) == len(state):

            for i in range(len(obj)):

                if not self._has_instance(state[i]):

                    obj[i] = self._get_pure(state[i])

                elif isinstance(state[i], tuple):

                    obj[i] = self._do_tuple(state[i])

                else:

                    self._do_object(obj[i], state[i])

        else:

            raise StateSetterError(

                'Cannot set state of list of incorrect size.'

            )



    def _do_dict(self, obj, state):

        for key, value in list(state.items()):

            if not self._has_instance(value):

                obj[key] = self._get_pure(value)

            elif isinstance(value, tuple):

                obj[key] = self._do_tuple(value)

            else:

                self._do_object(obj[key], value)





######################################################################

# Internal Utility functions.

######################################################################

def _get_file_read(f):

    if hasattr(f, 'read'):

        return f

    elif isinstance(f, str):

        return open(f, 'rb')

    else:

        raise TypeError('Given object is neither a file or String')



def _get_file_write(f):

    if hasattr(f, 'write'):

        return f

    elif isinstance(f, str):

        return open(f, 'wb')

    else:

        raise TypeError('Given object is neither a file or String')





######################################################################

# Utility functions.

######################################################################

def dump(value, file):

    """Pickles the state of the object (`value`) into the passed file

    (or file name).

    """

    f = _get_file_write(file)

    try:

        StatePickler().dump(value, f)

    finally:

        f.flush()

        if f is not file:

            f.close()





def dumps(value):

    """Pickles the state of the object (`value`) and returns a string.

    """

    return StatePickler().dumps(value)





def load_state(file):

    """Returns the state of an object loaded from the pickled data in

    the given file (or file name).

    """

    f = _get_file_read(file)

    try:

        state = StateUnpickler().load_state(f)

    finally:

        if f is not file:

            f.close()

    return state





def loads_state(string):

    """Returns the state of an object loaded from the pickled data

    in the given string.

    """

    return StateUnpickler().loads_state(string)





def get_state(obj):

    """Returns the state of the object (usually as a dictionary).  The

    returned state may be used directy to set the state of the object

    via `set_state`.

    """

    s = dumps(obj)

    return loads_state(s)





def set_state(obj, state, ignore=None, first=None, last=None):

    StateSetter().set(obj, state, ignore, first, last)

set_state.__doc__ = StateSetter.set.__doc__





def update_state(state):

    """Given the state of an object, this updates the state to the

    latest version using the handlers given in the version registry.

    The state is modified in-place.

    """

    version_registry.registry.update(state)





def create_instance(state):

    """Create an instance from the state if possible.

    """

    if (not isinstance(state, State)) and \

           ('class_name'  not in state.__metadata__):

        raise StateSetterError('No class information in state')

    metadata = state.__metadata__

    class_name = metadata.get('class_name')

    mod_name = metadata.get('module')

    if 'tvtk_classes' in mod_name:

        # FIXME: This sort of special-case is probably indicative of something

        # that needs more thought, plus it makes it tought to decide whether

        # this component depends on tvtk!

        from tvtk.api import tvtk

        return getattr(tvtk, class_name)()



    initargs = metadata['initargs']

    if initargs.has_instance:

        raise StateUnpicklerError('Cannot unpickle non-trivial initargs')



    __import__(mod_name, globals(), locals(), class_name)

    mod = sys.modules[mod_name]

    cls = getattr(mod, class_name)

    return cls(*initargs)


