#!/usr/bin/env python
import types
import marshal
import pickle

def dumpfunction(function):
    """
    Returns a serialized version of the given function.
    """
    return marshal.dumps((
            function.func_code,
            function.func_defaults,
            function.func_dict,
            function.func_doc,
            function.func_name))

def loadfunction(pickledfunction):
    """
    Returns a function loaded from data produced by the dumpfunction method
    of this module.
    """
    pairing = zip(
        ("func_code", "func_defaults", "func_dict", "func_doc", "func_name"),
        marshal.loads(pickledfunction))

    dummyfunction = lambda: None
    for attribute, value in pairing:
        setattr(dummyfunction, attribute, value)

    return dummyfunction


def dumpcontext(context, outstream=None):
    """
    Serializes the state of the environment passed into the function. A two
    element tuple is returned containing the serialized data and a list
    containing the names of the variables that could not be serialized.
    """
    codevat = dict()
    moduledict = dict()
    picklejar = dict()
    unpickled = []

    for variablename, variable in context.iteritems():
        if variablename.startswith("__"):
            continue

        if isinstance(variable, types.ModuleType):
            moduledict[variablename] = variable.__name__

        else:
            if (isinstance(variable, types.FunctionType)
                    and variable.__module__ == "__main__"):
                codevat[variablename] = dumpfunction(variable)
            else:
                try:
                    pickle.dumps(variable)
                    picklejar[variablename] = variable
                except Exception:
                    unpickled.append(variablename)

    container = (codevat, moduledict, pickle.dumps(picklejar))

    if outstream:
        pickle.dump(container, outstream)
        return unpickled

    return (pickle.dumps(container), unpickled)

def loadcontext(contextdump, target, raiseonerror=False):
    """
    Restores an interpreter environment using the data produced by the
    complimentary function dumpcontext.
    """
    if hasattr(contextdump, "read"):
        codevat, moduledict, picklejar = pickle.load(contextdump)
    else:
        codevat, moduledict, picklejar = pickle.loads(contextdump)

    module = "__%s__" % __name__
    exec("import pickle as __pickle__", target)
    exec("import %s as %s" % (__name__, module), target)

    for importname, modulename in moduledict.iteritems():
        try:
            exec("import %s as %s" % (importname, modulename), target)
        except:
            if raiseonerror:
                raise

    for name, dump in codevat.iteritems():
        try:
            exec("%s = %s.loadfunction(%r)" % (name, module, dump), target)
        except:
            if raiseonerror:
                raise

    exec("locals().update(__pickle__.loads(%r))" % picklejar, target)
    exec("del %s, __pickle__" % module, target)