"""
Extend pickle module to allow pickling of interpreter state
including any interactively defined functions and classes.

This module is not required for unpickling such pickle files.

>>> import savestate, pickle, __main__
>>> pickle.dump(__main__, open('savestate.pickle', 'wb'), 2)
"""

import sys, pickle, new

def save_code(self, obj):
    """ Save a code object by value """
    args = (
        obj.co_argcount, obj.co_nlocals, obj.co_stacksize, obj.co_flags, obj.co_code,
        obj.co_consts, obj.co_names, obj.co_varnames, obj.co_filename, obj.co_name,
        obj.co_firstlineno, obj.co_lnotab, obj.co_freevars, obj.co_cellvars
    )
    self.save_reduce(new.code, args, obj=obj)
pickle.Pickler.dispatch[new.code] = save_code

def save_function(self, obj):
    """ Save functions by value if they are defined interactively """
    if obj.__module__ == '__main__' or obj.func_name == '<lambda>':
        args = (obj.func_code, obj.func_globals, obj.func_name, obj.func_defaults, obj.func_closure)
        self.save_reduce(new.function, args, obj=obj)
    else:
        pickle.Pickler.save_global(self, obj)
pickle.Pickler.dispatch[new.function] = save_function

def save_global_byname(self, obj, modname, objname):
    """ Save obj as a global reference. Used for objects that pickle does not find correctly. """
    self.write('%s%s\n%s\n' % (pickle.GLOBAL, modname, objname))
    self.memoize(obj)

def save_module_dict(self, obj, main_dict=vars(__import__('__main__'))):
    """ Special-case __main__.__dict__. Useful for a function's func_globals member. """
    if obj is main_dict:
        save_global_byname(self, obj, '__main__', '__dict__')
    else:
        return pickle.Pickler.save_dict(self, obj)      # fallback to original 
pickle.Pickler.dispatch[dict] = save_module_dict

def save_classobj(self, obj):
    """ Save an interactively defined classic class object by value """
    if obj.__module__ == '__main__':
        args = (obj.__name__, obj.__bases__, obj.__dict__)
        self.save_reduce(new.classobj, args, obj=obj)
    else:
        pickle.Pickler.save_global(self, obj, name)
pickle.Pickler.dispatch[new.classobj] = save_classobj

def save_instancemethod(self, obj):
    """ Save an instancemethod object """
    # Instancemethods are re-created each time they are accessed so this will not be memoized
    args = (obj.im_func, obj.im_self, obj.im_class)
    self.save_reduce(new.instancemethod, args)
pickle.Pickler.dispatch[new.instancemethod] = save_instancemethod

def save_module(self, obj):
    """ Save modules by reference, except __main__ which also gets its contents saved by value """
    if obj.__name__ == '__main__':
        self.save_reduce(__import__, (obj.__name__,), obj=obj, state=vars(obj).copy())
    elif obj.__name__.count('.') == 0:
        self.save_reduce(__import__, (obj.__name__,), obj=obj)
    else:
        save_global_byname(self, obj, *obj.__name__.rsplit('.', 1))
pickle.Pickler.dispatch[new.module] = save_module

def save_type(self, obj):
    if getattr(new, obj.__name__, None) is obj:
        # Types in 'new' module claim their module is '__builtin__' but are not actually there
        save_global_byname(self, obj, 'new', obj.__name__)
    elif obj.__module__ == '__main__':
        # Types in __main__ are saved by value

        # Make sure we have a reference to type.__new__        
        if id(type.__new__) not in self.memo:
            self.save_reduce(getattr, (type, '__new__'), obj=type.__new__)
            self.write(pickle.POP)

        # Copy dictproxy to real dict
        d = dict(obj.__dict__)
        # Clean up unpickleable descriptors added by Python
        d.pop('__dict__', None)
        d.pop('__weakref__', None)
        
        args = (type(obj), obj.__name__, obj.__bases__, d)
        self.save_reduce(type.__new__, args, obj=obj)
    else:
        # Fallback to default behavior: save by reference
        pickle.Pickler.save_global(self, obj)
pickle.Pickler.dispatch[type] = save_type