Advanced Uses of Decorators

Let's jump right in!

>>> def trace(f):
...     "Decorate a function to print a message before and after execution."
...     def traced(*args, **kw):
...         "Print message before and after a function call."
...         print("Entering", f.__name__)
...         result = f(*args, **kw)
...         print("Leaving", f.__name__)
...         return result
...     return traced
...
>>> @trace
... def myfunc(x, a=None):
...     "Simply prints a message and arguments."
...     print("Inside myfunc")
...     print("x:", x, "a:", a)
...
>>> myfunc("ONE", "TWO")
Entering myfunc
Inside myfunc
x: ONE a: TWO
Leaving myfunc
>>>

In the example above, the trace function is a decorator. That means that it takes a single argument (which is normally the function being decorated). Internally, it defines a function traced() that prints out a line of text, calls the decorated function with whatever arguments it was called with itself, prints out another line of text and then returns the result obtained from the decorated function. Then, trace returns the function it has just defined.

This means that you can apply trace() to any function, and the result will do just what the original function did as well as printing out a line before and after the call to the decorated function. This is how most decorators work (although as always there are some smart people who have found non-standard ways to use decorators that were not originally intended by the specification). That's why you often see the internal function written to accept any combination of positional and keyword arguments—it means that the decorator can be applied to any function, no matter what its signature.

Remember, the decorator syntax is really just an abbreviation; it doesn't do anything that you couldn't do without the syntax. When you write @trace before the definition for myfunc(), it's exactly equivalent to writing myfunc = trace(myfunc) after the function definition. The syntax was added because with longer function definitions it was often difficult to notice the reassignment to the name when it followed the function definition. The feature was restricted to functions when it was originally introduced, but now you can also decorate classes. While this is a little bit more complicated than decorating functions, it does have its uses.

Because the above decorator defines a function that contains a call to the decorated function as a part of its code (traced() in the example above), we say that the decorator wraps the decorated function. This has certain unfortunate side effects: mostly, the name of the function appears to change to the name of the wrapper function from inside the decorator, and the docstring is that of the wrapper.

>>> trace.__name__          # undecorated
'trace'
>>> myfunc.__name__         # decorated
'traced'
>>> myfunc.__doc__
'Print message before and after a function call.'
>>>

Fortunately, this issue can be handled using the wraps decorator from the functools library. This is provided precisely to ensure that decorated functions continue to "look like themselves." Until you get the hang of using it, however, it seems a little weird because it means you end up using a decorator on the wrapper function inside your decorator! But honestly, it isn't difficult.

>>> from functools import wraps
>>> def simpledec(f):
...     "A really simple decorator to demonstrate functools.wraps"
...     @wraps(f)
...     def wrapper(arg):
...         print("Calling f with arg", arg)
...         return f(arg)
...     return wrapper
...
>>> @simpledec
... def f(x):
...     "Simply prints its argument."
...     print("Inside f, arg is", x)
...
>>> f("Hello")
Calling f with arg Hello
Inside f, arg is Hello
>>> f.__name__
'f'
>>> f.__doc__
'Simply prints its argument.'
>>>

While decorators are usually functions, they don't need to be—any callable can be used as a decorator. This means that you could use a class as a decorator, and when the decoration takes place the class's __init__() method is called with the object to be decorated (whether it's a function or a class: note that a decorator is typically designed to decorate either functions or classes but not both because they are fairly different in nature).

If you want to decorate a function with a class, remember that calling a class calls its __init__() method, and returns an instance of the class. As always, the first argument to __init__() is self, the newly created instance, so in this case the function that the interpreter passes to the decorator will end up as the second argument to __init__(). Since calling the class creates an instance, and since normally you want to be able to call the decorated function, the classes you use as decorators should define a __call__() method, which will then be called when the decorated function is called.

>>> class ctrace:
...     def __init__(self, f):
...         "__init__ records the passed function for later use in __call__()."
...         self.__doc__ = f.__doc__
...         self.__name__ = f.__name__
...         self.f = f
...     def __call__(self, *args, **kw):
...         "Prints a trace line before calling the wrapped function."
...         print("Called", self.f.__name__)
...         return self.f(*args, **kw)
...
>>> @ctrace
... def simple(x):
...     "Just prints arg and returns it."
...     print("simple called with", x)
...     return x
...
>>> simple("walking")
Called simple
simple called with walking
'walking'
>>> simple.__name__
'simple'
>>> simple.__doc__
'Just prints arg and returns it.'
>>>

By the time the decorator is called, the simple() function has already been compiled, and it is passed to the decorator's __init__() method, where it is stored as an instance variable. To make sure the decorated function retains its name and docstring, those attributes of the function are also copied into instance variables with the same names.

Up until now, we have decorated functions, but once the feature was introduced into Python, it was only a matter of time before it was extended to classes. So now you can decorate classes in just the same way as functions. The principle is exactly the same: the decorator receives a class as an argument, and (usually) returns a class. Because classes are more complicated than functions you will find it most convenient to modify the class in place and return the modified class as the result of the decorator.

To demonstrate this, suppose that you want to be able to have each of the methods of a class print out a trace call during debugging. You could simply apply the trace decorator above to each method, but that would mean extensive editing for a large class when you wanted to switch the debugging off. It is simpler for programmers to use a class decorator, so we might well accept a slightly higher level of complexity in the decorator to avoid the editing. Once the interpreter has processed the class definition, it calls the decorator with the class as its argument, and the decorator can then either create a new class (which is fairly difficult) or modify the class and return it.

Since the interactive session has already defined a simple tracing function, we'll use that to wrap each of the methods in our decorated class. Finding the methods is not as easy as you might imagine. It involves looking through the class's __dict__ and finding callable items whose names do not begin and end with "__" (it's best not to mess with the "magic" methods). Once such an item is found, it is wrapped with the trace() function and replaced in the class __dict__.

>>> def callable(o):
...     return hasattr(o, "__call__")
...
>>> def mtrace(cls):
...     for key, val in cls.__dict__.items():
...         if key.startswith("__") and key.endswith("__") \
...                     or not callable(val):
...             continue
...         setattr(cls, key, trace(val))
...         print("Wrapped", key)
...     return cls
...
>>> @mtrace
... class dull:
...     def method1(self, arg):
...         print("Method 1 called with arg", arg)
...     def method2(self, arg):
...         print("Method 2 called with arg", arg)
...
Wrapped method1
Wrapped method2
>>> d = dull()
>>> d.method1("Hello")
Entering method1
Method 1 called with arg Hello
Leaving method1
>>> d.method2("Goodbye")
Entering method2
Method 2 called with arg Goodbye
Leaving method2
>>>

As you can see, when you call method1() and method2(), they print out the standard "before and after" trace lines, because they are now wrapped by the trace() function.

Sometimes you don't want to wrap the function: instead you want to alter it in some other way, such as adding attributes (yes, you can add attributes to functions the same way as you can to most of the other objects in Python). In that case, the decorator simply returns the function that is passed in as an argument, having modified the function in whatever ways it needs to. So next we'll write a decorator that flags a function as part of a framework by adding a "framework" attribute.

>>> def framework(f):
...     f.framework = True
...     f.author = "Myself"
...     return f
...
>>> @framework
... def somefunc(x):
...     pass
...
>>> somefunc.framework
True
>>> somefunc.author
'Myself'
>>>

Note that the decorator does still return a function, but since there is no need to wrap the decorated function it simply returns the function that it was passed (now resplendent with new attributes). Since this avoids a second function call, it will be slightly quicker and there is no need to use functools.wraps because the function is not being wrapped.

Python includes two built-in functions that are intended for use in decorating methods. The staticmethod() function modifies a method so that the special behavior of providing the instance as an implicit first argument is no longer applied. In fact, the method can be called on either an instance or the class itself, and it will receive only the arguments explicitly provided to the call. It becomes a static method. You can think of static methods as being functions that don't need any information from either their class or their instance, so they do not need a reference to it. Such functions are relatively infrequently seen in the wild.

If you want to write a method that relies on data from the class (class variables are a common way to share data among the various instances of the class) but does not need any data from the specific instance, you should decorate the method with the classmethod() function to create a class method. Like static methods, class methods can be called on either the class or an instance of the class. The difference is that the calls to a class method do receive an implicit first argument. Unlike a standard method call, though, this first argument is the class that the method was defined on rather than the instance it was called on. The conventional name for this argument is cls, which makes it more obvious that you are dealing with a class method.

You may well ask what static and class methods are for—why use them when we already have standard methods that are perfectly satisfactory for most purposes? Why not just use functions instead of static methods, since no additional arguments are provided? The answer to this question lies in the fact that these functions are methods of a class, and so will be inherited (and can be overridden or extended) by any subclasses you may define. Further, the instances of the class can reference class variables rather than using a global—this is always safer because there is no guarantee, when your code lands in someone else's program, that their code isn't using the same global name for some other purpose. It is difficult to think of any example where the use of a classmethod would be absolutely required, but sometimes it can simplify your design a little.

A typical application for class methods has each of the instances using configuration data that is common to all, and saved in the class. If you provide methods to alter the configuration data (for example, changing the frequency a wireless transmitter works on, or changing the function that the instances call to allocate resources), they do not need to reference any of the instances, so a class method would be ideal.

Sometimes you want to write a decorator that takes parameters. Remember, though, that the decorator syntax requires a callable that takes precisely one argument (the class or function to be decorated). So if you want to parameterize a decorator, you have to do so "at one remove"—the function that takes the arguments has to return a function that takes one argument and returns the decorated object. This can be a little brain-twisting, so an example may help. Or, it may just make your head explode!

Suppose that you wanted to have your program record the number of calls that are made to each of several different types of function. When you define a function, you want to give a parameter to the decorator to specify the classification of the decorated function.

@countable('special')
def f(...):
    ...

In other words, @countable('special') has to return a function that is a conventional decorator—it takes a single function as an argument and returns the decorated version of the function as its result. This means that we need to nest functions three levels deep! We will use a global variable to store a dict, and the different function-type strings will be the keys. Here we go!

>>> counts = {}
>>> def countable(ftype):
...     "Returns a decorator that counts each call of a function against ftype."
...     def decorator(f):
...         "Decorates a function and to count each call."
...         def wrapper(*args, **kw):
...             "Counts every call as being of the given type."
...             try:
...                 counts[ftype] += 1
...             except KeyError:
...                 counts[ftype] = 1
...             return f(*args, **kw)
...         return wrapper
...     return decorator
...
>>> @countable("short")
... def f1(a, b=None):
...     print("f1 called with", a, b)
...
>>> @countable("f2")
... def f2():
...     print("f2 called")
...
>>> @countable("short")
... def f3(*args, **kw):
...     print("f3 called:", args, kw)
...
>>> for i in range(10):
...     f1(1)
...     f2()
...     f3(i, i*i, a=i)
...
f1 called with 1 None
f2 called
f3 called: (0, 0) {'a': 0}
f1 called with 1 None
f2 called
f3 called: (1, 1) {'a': 1}
f1 called with 1 None
f2 called
f3 called: (2, 4) {'a': 2}
f1 called with 1 None
f2 called
f3 called: (3, 9) {'a': 3}
f1 called with 1 None
f2 called
f3 called: (4, 16) {'a': 4}
f1 called with 1 None
f2 called
f3 called: (5, 25) {'a': 5}
f1 called with 1 None
f2 called
f3 called: (6, 36) {'a': 6}
f1 called with 1 None
f2 called
f3 called: (7, 49) {'a': 7}
f1 called with 1 None
f2 called
f3 called: (8, 64) {'a': 8}
f1 called with 1 None
f2 called
f3 called: (9, 81) {'a': 9}
>>> for k in sorted(counts.keys()):
...     print(k, ":", counts[k])
...
f2 : 10
short : 20
>>>

As you can see, f1 and f3 are classified as "short", while f2 is classified as "f2." Every time a @countable function is called, one is added to the count for its category. There were 30 function calls in all, 20 to category "short" (f1 and f3). Calling countable() returns a decorator whose action is to add one to the count identified by its argument. Your code defines a function (countable()) that defines a function (decorator(), which is a decorator, that defines a function (wrapper) that wraps the function f provided as an argument to decorator, which was produced by calling countable. This is probably about as far as anyone wants to go with decorators (and a little bit further than most).