Wednesday, February 9, 2011

Creating Custom Function Decorator in Python

Python provides many built-in function decorators, such as staticmethod, abstractmethod, etc. Creating our own function decorator is pretty easy. What we need to do is to overload the __call__ function. Here is how to do it.

class Log:
    def __init__(self, func):
        self.func = func
    def __call__(self, *args):
        print "Before", self.func.__name__ + "()"
        self.func(*args)
        print "After", self.func.__name__ + "()"

@Log
def hello():
    print "Hello"
    
if __name__ == "__main__":
    hello()

This is equivalent to:
class Log:
    def __init__(self, func):
        self.func = func
    def __call__(self, *args):
        print "Before", self.func.__name__ + "()"
        self.func(*args)
        print "After", self.func.__name__ + "()"

def hello():
    print "Hello"
hello = Log(hello)
    
if __name__ == "__main__":
    hello()

The output is
Before hello()
Hello
After hello()

No comments:

Post a Comment