def wrap(fn):
    if fn is None or fn.__class__ is _StackContextWrapper:
        return fn
    # functools.wraps doesn't appear to work on functools.partial objects
    #@functools.wraps(fn)

    def wrapped(callback, contexts, *args, **kwargs):
        if contexts is _state.contexts or not contexts:
            callback(*args, **kwargs)
            return
        if not _state.contexts:
            new_contexts = [cls(arg) for (cls, arg) in contexts]
        # If we're moving down the stack, _state.contexts is a prefix
        # of contexts.  For each element of contexts not in that prefix,
        # create a new StackContext object.
        # If we're moving up the stack (or to an entirely different stack),
        # _state.contexts will have elements not in contexts.  Use
        # NullContext to clear the state and then recreate from contexts.
        elif (len(_state.contexts) > len(contexts) or
            any(a[1] is not b[1]
                for a, b in izip(_state.contexts, contexts))):
            # contexts have been removed or changed, so start over
            new_contexts = ([NullContext()] +
                            [cls(arg) for (cls,arg) in contexts])
        else:
            new_contexts = [cls(arg)
                            for (cls, arg) in contexts[len(_state.contexts):]]
        if len(new_contexts) > 1:
            with _nested(*new_contexts):
                callback(*args, **kwargs)
        elif new_contexts:
            with new_contexts[0]:
                callback(*args, **kwargs)
        else:
            callback(*args, **kwargs)
    if _state.contexts:
        return _StackContextWrapper(wrapped, fn, _state.contexts)
    else:
        return _StackContextWrapper(fn)
Returns a callable object that will restore the current StackContext
dataset/ETHPy150Open felipecruz/zmqpy/zmqpy/eventloop/stack_context.py/wrap