gistfile1.py

def add_method_wrapper(cls, wrapper_func):
    """ Substitute the function named "wrapperfunc.__name__" at class
    "cls" with a function that wraps the call to the original function.
    Return an undo function which can be called to reset the class to use
    the old method again.

    wrapper_func is called with the same arguments as the method
    it wraps and its result is used as a wrap_controller for
    calling the original function.
    """
    name = wrapper_func.__name__
    oldcall = getattr(cls, name)
    def wrap_exec(*args, **kwargs):
        gen = wrapper_func(*args, **kwargs)
        return wrapped_call(gen, lambda: oldcall(*args, **kwargs))

    setattr(cls, name, wrap_exec)
    return lambda: setattr(cls, name, oldcall)


def wrapped_call(wrap_controller, func):
    """ Wrap calling to a function with a generator.  The first yield
    will trigger calling the function and receive an according CallOutcome
    object representing an exception or a result.
    """
    next(wrap_controller)   # first yield
    call_outcome = CallOutcome(func)
    try:
        res = wrap_controller.send(call_outcome)
    except StopIteration:
        pass
    if call_outcome.excinfo is None:
        return call_outcome.res
    else:
        py.builtin._reraise(*call_outcome.excinfo)


class CallOutcome:
    res = excinfo = None
    def __init__(self, func):
        try:
            self.res = func()
        except Exception:
            self.excinfo = sys.exc_info()

    def force_result(self, res):
        self.excinfo = None
        self.res = res