monkeypatch.py

def monkeypatch(mod, what):
    """
    Patch function named <what> from module <mod> to run the decorated function after <what> completes.

    Eg::

        @monkeypatch(os, 'forkpty')
        def patched_forkpty(pid_fd):
            pid, fd = pid_fd
            if not pid:
                fixup_stuff()
            return pid, fd
    """

    old = getattr(mod, what)
    def decorator(func):
        def patch():
            ret = old()
            try:
                func(ret)
            except:
                traceback.print_exc()
            return ret
        setattr(mod, what, patch)
    return decorator