#!/usr/bin/python

class Stub(object):
    """
    A simple class to restore/replace objects. 
    It's not complete, but it's enough for many cases.
    """

    def __init__(self):
        self.funcs = {}

    def replace(self, module, function_name, new_function):
        self.funcs.setdefault(module, {})[function_name] = getattr(module, function_name)
        setattr(module, function_name, new_function)

    def restore(self, module, function_name):
        assert getattr(module, function_name)

        setattr(module, function_name, self.funcs[module][function_name])

# Demo for Stub.
class MyClass(object):
    def hello(self, msg):
        print msg

def new_hello(self, msg):
    print msg + " (by replaced method)"

obj = MyClass()
obj.hello('Testing rocks!')  # Testing rocks!
mystub = stub.Stub()
mystub.replace(MyClass, 'hello', new_hello)
obj.hello('Testing rocks!')  # Testing rocks! (by replaced method)
mystub.restore(MyClass, 'hello')
obj.hello('Testing rocks!')  # Testing rocks!