1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

#!/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!