Python Decorator class with arguments

Question

I have a simple Python class that I want to use to add named hooks to a program I am writing. I try to run the below code and I get the following output.

---

Code:

hooks = {}

class hook(object):
    def __init__(self, f, hook):
        if hook not in hooks:
            hooks[hook] = []

        hooks[hook].append({"module": f.__module__, "func": f})
        self.f = f

    def __call__(self, *args):
        f(*args)

@hook("test")
def testHook():
    print "hi"

---

Output:

Traceback (most recent call last):                       
  File "<stdin>", line 1, in <module>                    
TypeError: __init__() takes exactly 3 arguments (2 given)

How can I fix this? I am using Python 2.7

Answer 1

Your decorator syntax:

@hook("test")
def testHook():
    # ...

translates to:

def testHook():
    # ...

testHook = hook("test")(testHook)

so it is given just one argument. You need to restructure your code to produce your class decorator as a return value of hook().

The following would work:

def hook(hookname):
    class HookDecorator(object):
        def __init__(self, f):
            if hookname not in hooks:
                hooks[hookname] = []

            hooks[hook].append({"module": f.__module__, "func": f})
            self.f = f

        def __call__(self, *args):
            return self.f(*args)

    return HookDecorator

where the __call__ uses self.f, not the global f, and returns whatever the decorated function produced. I'm assuming that hooks is a global dictionary somewhere.

There is little point in making this a class however; you use very little state. You could just as well use a wrapper function here:

from functools import wraps

def hook(hookname):
    def decorator(f):
        if hookname not in hooks:
            hooks[hookname] = []

        hooks[hookname].append({"module": f.__module__, "func": f})

        @wraps(f)
        def wrapper(*args):
            return f(*args)

        return wrapper

    return decorator