Writing/The Unbearable Lightness of Coding
Chapter 8 · 37 min read
Decorators: Mastering Function Objects
In this chapter, we're going to talk about decorators: a particularly powerful design pattern that relies on a thorough understanding of functions, which we now claim to possess.
We'll cover the syntax, as always; but also the evolution of this idea, how first- and second-order decorators work, and some important nuances to consider when writing a decorator.
This time, though, we're not going to start with the syntax, because it doesn't really make sense before we can articulate what is it that we're trying to do. We talked a lot about functions in these last few chapters, but we only touched on the different ways in which they can be composed and applied to each other — so to build some motivation, let's revisit the topic again.
Functions encapsulate input-output relations, but some of these relations are more meta than others. A math function that calculates a formula takes some numbers, and returns a number — so it's a relation that operates on data. But in more complex cases, we might be operating on entire behaviors: when saving some text to a file, for example, we might also want to compress it, and/or encrypt it, and/or back it up, and who knows what else. In fact, let's do that — and to avoid irrelevant implementation details, I'll just print what happens, symbolically reassign the value, and pretend that the desired changes or side-effects took place:
>>> def compress(text):
... print('compressing', text)
... compressed_text = text
... return compressed_text
>>> def encrypt(text):
... print('encrypting', text)
... encrypted_text = text
... return encrypted_text
>>> def back_up(path):
... print('backing up', path)
The point is, while all these relations operate on data, a flow that combines them dynamically would be a relation that operates on relations. In our case, the entry point is going to be save, which receives some text, generates a file for it, writes it, and returns its path:
>>> def save(text):
... path = 'path'
... print('saving', text, 'to', path)
... return path
Configuring Behaviors
So how do we equip save with our additional behaviors, in a way that lets us toggle them on and off? One way to do it, if we think of arguments in terms of data, would be to pass, well... additional data — booleans, for instance, to specify which behaviors are on:
>>> def save(text, compression=False, encryption=False, backup=False):
... if compression:
... text = compress(text)
... if encryption:
... text = encrypt(text)
... path = 'path'
... print('saving', text, 'to', path)
... if backup:
... back_up(path)
... return path
This way we can choose to have compression and backup, but no encryption, like this:
>>> path = save('text', compression=True, backup=True)
compressing text
saving text to path
backing up path
It works, but introduces some problems. When designing any system, we should ask ourselves how it is going to change over time. If it's not going to, then whatever we do will be good enough, as long as it works; but if change is imminent, we should plan ahead and ensure that our system is as easy and safe to maintain and extend as we can make it.
For example, some behaviors might require extra configuration, like what compression or encryption algorithm to apply, or what backup directory to use. Luckily, we can support it easily enough — instead of plain booleans, we can pass in entire dictionaries, and simply forward them to each component:
>>> def compress(text, config):
... if 'algorithm' in config:
... print('using', config['algorithm'], 'compression')
... # same as before
>>> def save(text, compression=None, ...):
... if compression is not None:
... text = compress(text, compression)
... # same as before
>>> path = save('text', compression={'algorithm': 'gzip'})
using gzip compression
compressing text
saving text to path
We can even support the case in which compression=True, indicating that it should be on, but with default behavior. To do so, we can rewrite the component functions to receive their configurations as actual arguments, and pass those (or an empty dictionary, in case of True) by keyword:
>>> def compress(text, algorithm=None):
... if algorithm:
... print('using', algorithm, 'compression')
... # same as before
>>> def save(text, compression=None, ...):
... if compression:
... config = {} if compression is True else compression
... text = compress(text, **config)
... # same as before
>>> path = save('text', compression={'algorithm': 'gzip'})
using gzip compression
compressing text
saving text to path
>>> path = save('text', compression=True)
compressing text
saving text to path
Extending Behaviors
What about adding or removing behaviors? Suppose we want to normalize some funny characters to more standard ones, or to change the permissions of the file after it's created — this would require restructuring the main function, which is less than ideal, since over time it will grow bigger and messier.
And what if in some scenarios, we want to save the text differently, like storing it in a database and returning its ID, or uploading it to a webserver and returning its URL? Assuming we want to have the same manipulations available, do we put all the various saving logics in the same function, or fork it into separate ones, repeating all the non-saving logics in each?
It seems we need a different approach — one in which our various logics are decoupled, so we can compose them with more flexibility. In Python, it's easy to imagine a list of functions, which we can iterate over and invoke:
>>> def save(text, transforms=None):
... for transform in transforms or []:
... text = transform(text)
... path = 'path'
... print('saving', text, 'to', path)
... return path
>>> path = save('text', [compress, back_up])
compressing text
backing up text # ???
saving text to path
The problem is that some manipulations, like compression and encryption, operate on the text before it's written, while others, like backup, operate on the path after the file is properly saved. We can manage the two cases in two separate lists:
>>> def save(text, before=None, after=None):
... for transform in before or []:
... text = transform(text)
... path = 'path'
... print('saving', text, 'to', path)
... for transform in after or []:
... transform(path)
... return path
>>> path = save('text', before=[compress], after=[backup])
compressing text
saving text to path
backing up path # ✓
But then, how do we add configuration back in the mix? I guess we can make it a list of pairs, grouping each function with its arguments and passing them on like before:
>>> def save(text, before=None, after=None):
... for transform, kwargs in before or []:
... text = transform(text, **kwargs)
... path = 'path'
... print('saving', text, 'to', path)
... for transform, kwargs in after or []:
... transform(path, **kwargs)
... return path
>>> path = save(
... text = 'text',
... before = [(compress, {'algorithm': 'gzip'})],
... after = [(backup, {})],
... )
using gzip compression
compressing text
saving text to path
backing up path
This successfully addresses all the problems that we've discussed — but it ends up pretty clunky and fragile; flexible, sure, but hardly good design.
Maybe a better way to organize everything would be with object-oriented programming: we can define a Saver class and a bunch of behavior classes, like Compression, Encryption and BackUp, then use inheritance or composition to combine them in all sorts of ways. It has its strengths and weaknesses, and we'll talk more about it once we've covered classes; for now, let's examine a more functional approach.
Composing Behaviors
To start, let's revert the save function to its simplest form:
>>> def save(text):
... path = 'path'
... print('saving', text, 'to', path)
... return path
As long as we keep it focused on its task, we'll have no problem supporting different saving logics, from databases to webservers, in small, separate and similarly dedicated functions. The only question is how to equip our various components on top — to compose any one saving behavior with all the other manipulations without repeating them.
Imagine, for a moment, that we could implement compress, encrypt and back_up in the following way: they would have the same signature as save and receive the same arguments, but also have access to a magical proceed function. This function, once called, will carry out "the rest of the saving process", whatever it may be. In this case, our implementations would look like this:
>>> def compress(text):
... print('compressing', text)
... compressed_text = text
... return proceed(compressed_text)
>>> def encrypt(text):
... print('encrypting', text)
... encrypted_text = text
... return proceed(encrypted_text)
>>> def back_up(text):
... path = proceed(text)
... print('backing up', path)
... return path
Note how compress and encrypt work on the text first, then proceed on the manipulated version; and how backup actually starts by calling proceed, and only does its thing afterwards, once it knows which path to back up.
All that remains is to figure out where to get that magical proceed function. It can't be a static, global object, because it's constantly changing — in different saving flows, or even different moments of the same flow, "do the rest of the saving process" means different things. In other words, it's variable; so we can introduce it by way of an argument, like this:
>>> def compress(proceed):
... def do_compress(text):
... print('compressing', text)
... compressed_text = text
... return proceed(compressed_text)
... return do_compress
>>> def encrypt(proceed):
... def do_encrypt(text):
... print('encrypting', text)
... encrypted_text = text
... return proceed(encrypted_text)
... return do_encrypt
>>> def back_up(proceed):
... def do_back_up(text):
... path = proceed(text)
... print('backing up', path)
... return path
... return do_back_up
I know, it's weird, but bear with me. We wrap each of the original components (changing their name to do_*) in a factory function. Instead of carrying out the manipulation immediately, this function returns a new function that will do so when called — and in which proceed actually makes sense, because it's bound to a concrete argument.
To figure out what this argument might be, let's start with the simplest case. Suppose we want to save something with just one manipulation, like compression. In this scenario, proceed — that is, "the rest of the saving process" — is really just the save function itself, isn't it? We pass it in:
>>> compress(save)
<function do_compress() at 0x...>
… and it returns a function, right. But isn't this exactly the function we want? — do_compress, a flow that applies compression and then proceeds to save? If we invoke it on our desired argument:
>>> path = compress(save)('text')
compressing text
saving text to path
It seems to work! Let's throw in another component, like backing up. Now, for the proceed function — again, "the rest of the saving process" — we can use do_compress; it successfully combined both save and compress, after all, and that's exactly what we want to further combine with back_up:
>>> back_up(compress(save))
<function do_back_up() at 0x...>
>>> path = back_up(compress(save))('text')
compressing text
saving text to path
backing up path
And all three:
>>> path = back_up(encrypt(compress(save)))('text')
compressing text
encrypting text
saving text to path
backing up path
Quite unusual, but pretty cool, isn't it? This design pattern is called the decorator, since we're effectively taking some initial flow — in our case, the save function — and decorating it with additional logics. Each decorator takes "the function so far", which we called proceed, and does something before and/or after calling it — maybe augments some arguments, like compress and encrypt do to the text, or follows up on the return value, like backup does with the path. The only requirement is that the signatures of the functions are the same; this way, we can wrap them in each other, and each one can successfully invoke the next — calling all the way into the innermost original flow, and then unfolding and returning all the way back.
To make it support configuration like before, all we have to do is add the desired parameters to the outer functions, like so:
>>> def compress(proceed, algorithm=None):
... def do_compress(text):
... if algorithm:
... print('using', algorithm, 'compression')
... print('compressing', text)
... compressed_text = text
... return proceed(compressed_text)
... return do_compress
And then embed them in the right place as part of our composition:
>>> back_up(compress(save, algorithm='gzip'))('text')
using gzip compression
compressing text
saving text to path
backing up path
It's a cute idea, and a useful trick. Unfortunately, if you read about this design pattern elsewhere, chances are you'll find a very different explanation — and from my experience, one that's doesn't map to Python very well. It's the case with many "classical" design patterns: something about them feels awkward and inadequate when applied to modern programming languages — a phenomenon that we'll try to understand and ameliorate in a later book. For now, just think of it as a nifty functional programming technique, and hopefully one that makes sense: all we did was find a way to decouple different logics into clean, separate functions, then compose them dynamically, mixing and matching them however we want on demand. In a system that keeps changing, as most systems tend to do over time, this approach is infinitely more maintainable than clumping everything together in a big, complex monolith — especially when it becomes too much and has to be duplicated.
Decoration Syntax
In fact, this turns out to be such a useful practice that Python has special syntax for exactly such compositions. To give a silly example, suppose we have a function that increments a number:
>>> def inc(x):
... return x + 1
And for some odd reason, we want to double its result — but we want to do it in a composable way, like before, by applying an additional layer of "doubling" behavior without touching the original function.
So, we define a double decorator, which takes the original function — what we called proceed, but will simply denote with f from now on — and returns a new function, often called a wrapper. This wrapper can call the original function, since it's part of its scope; except it also multiplies its return value by two:
>>> def double(f):
... def wrapper(x):
... return 2 * f(x)
... return wrapper
Now, if we double(inc), we get a composite function; and when we call it, we get a doubled result:
>>> double(inc)(1)
4
Except the naming conventions, it's really very similar to what we've already seen. And if we'd want this composition to be permanent, we can even replace inc with its so called decorated version entirely:
>>> inc = double(inc)
>>> inc(1)
4
This way, it always returns a doubled result. That last step — replacing inc with the decorated version of itself — can be done in Python upon definition. By using the asperand or "at" sign (@) with the decorator right above a function definition, we effectively hijack it: as soon as it's defined it gets passed into that decorator, and seamlessly swapped with its result for good:
>>> @double
... def inc(x):
... return x + 1
This way, inc comes into the world already doubled, and calling it on 1 returns 4 right from the onset:
>>> inc(1)
4
To apply this to our previous example, let's rewrite our components in a similar way:
>>> def compress(f):
... def wrapper(text):
... print('compressing', text)
... compressed_text = text
... return f(compressed_text)
... return wrapper
>>> def encrypt(f):
... def wrapper(text):
... print('encrypting', text)
... encrypted_text = text
... return f(encrypted_text)
... return wrapper
>>> def back_up(f):
... def wrapper(text):
... path = f(text)
... print('backing up', path)
... return path
... return wrapper
With these decorators, we can define the save function already compressed, encrypted and backed-up like so:
>>> @back_up
... @encrypt
... @compress
... def save(text):
... path = 'path'
... print('saving', text, 'to', path)
... return path
>>> path = save('text')
compressing text
encrypting text
saving text to path
backing up path
Or, if we decide that we don't want encryption after all, we can simply remove (or even comment out) that particular decoration:
>>> @back_up
... # @encrypt
... @compress
>>> def save(text):
... path = 'path'
... print('saving', text, 'to', path)
... return path
>>> path = save('text')
compressing text
saving text to path
backing up path
This notation lets us literally stack behaviors on top of each other — taking one function and returning a decorated, enhanced version of it, composing all the different functionalities together.
Permanance and Ambivalence
But if we're honest, this use case is a bit too obscure to justify new syntax, isn't it? If we're going through the trouble of defining functions in a composable way, it doesn't really matter if we compose them with a funky declarative line or with an explicit invocation. The code above is equivalent to a regular definition, followed by:
>>> save = back_up(compress(save))
What's more, it equips save with some components and then overwrites it, so we can't compose it differently on demand, like we did before — in a way, it's actually worse.
So what's the catch? Looking at a component again:
>>> def compress(f):
... def wrapper(text):
... print('compressing', text)
... compressed_text = text
... return f(compressed_text)
... return wrapper
We might notice that it doesn't really have anything to do with save at this point, does it? We can apply it to any other function that receives a string:
>>> @compress
... def upload(data):
... print('uploading', data)
>>> upload('data')
compressing data
uploading data
In other words, we can concieve of generic decorators, ones that can be useful to all sorts of functions — and then apply them with this clean, elegant syntax, simply by stacking them on top of a definition. Imagine adding an @optimized decoration, and having your function run faster; or adding a @synchronized decoration, and having it run in a thread-safe manner; or adding an @admin decoration, so it only works for administrative users; or adding a @trace decoration, so it prints out some logs whenever it runs; the possibilities are endless.
Generic Decorators: a New Idea
Before we dive into them, though, I'd like to tie everything together. Like I said, there is some confusion about decorators — especially for people who are coming from other languages, or reading about this design pattern elsewhere and trying to carry it over to Python.
The name and the idea of the decorator is exactly what we saw: it lets us compose different behaviors dynamically. But in most languages, it faces a severe limitation — in order for the functions to be composable, their signatures have to be the same. compress, encrypt, and back_up all had to have the same parameters and return value as save in order to call each other interchangeably. As a result, such decorators tend to be pretty specific — the point is not to reuse them in other scenarios, which will probably involve different signatures, but to compose this particular scenario in any way we want. From this angle, the whole notion of a generic decorator is nonsensical.
In Python, however, this structure has morphed and developed into something else, in an evolution that is rarely acknowledged. Looking at our double decorator again:
>>> def double(f):
... def wrapper(x):
... return 2 * f(x)
... return wrapper
It, too, was actually quite specific to inc, in that their signatures are the same. If we were to double something else:
>>> @double
... def add(x, y):
... return x + y
>>> add(1, 2)
TypeError: wrapper() takes 1 positional argument but 2 were given
It wouldn't work, since add gets replaced with wrapper, which expects a single x. Nor does a single x work, for that matter:
>>> add(1)
TypeError: add() missing 1 required positional argument: 'y'
Because the wrapper attempts to invoke the original function on this sole x, whereas add also expects a y. In other words, the signatures don't match, the decorator isn't generic, and the whole thing doesn't work.
Still, there's nothing inherently inc-specific in the concept of "doubling the result", is there? At least in theory, we could apply a similar doubling decoration to add. To do that, we can either rewrite the wrapper to accept x and y, in which case it will only be useful for that signature; or, we can use the omnisignature with perfect forwarding:
>>> def double(f):
... def wrapper(*args, **kwargs):
... return 2 * f(*args, **kwargs)
... return wrapper
This way, double creates a wrapper that accepts any arguments, passes them on to the original function precisely in the same way — but also does its thing, doubling the result. It means we can now apply the same decorator to inc, add, and anything else, really:
>>> @double
... def inc(x):
... return x + 1
>>> inc(1)
4
>>> @double
... def add(x, y):
... return x + y
>>> add(1, 2)
6
And it works in all cases; a generic decorator indeed.
So in Python, decorators are not just about composing particular behaviors with limited signatures; we can actually use the same pattern to compose generic behaviors with any signature. All the examples we mentioned before — optimization, synchronization, checking permissions, tracing logs — these are known as cross-cutting concerns; features that can be relevant to any function, with any signature, which is definitely not what the original decorator was about.
In fact, in the proposal to add decorators to Python, it's even mentioned that some people complained about the name, saying it's not consistent with the design pattern — but I hope I was able to explain how the two are related, at least; how the first one came about, and how the second one actually expanded it into a whole new realm of possibilities.
Example: Tracing
With that in mind, it's high time we implement an actual Python decorator; I have a feeling you weren't sufficiently impressed by the usefulness of my doubling example. Suppose we want to trace the execution of a function, so we can debug it. We can always add some print statements inside it — like, right after it's called, right before it ends — but that's something we'd have to maintain manually in every function that we want to trace; and if we're honest, it's completely unrelated to what the function actually does, which makes it a perfect candidate for a generic decorator.
All we need to do, given some function, is to replace it with a wrapper that prints something before we enter it, and something after we leave it:
>>> def trace(f):
... def wrapper(*args, **kwargs):
... print('enter', f.__name__)
... try:
... return f(*args, **kwargs)
... finally:
... print('leave', f.__name__)
... return wrapper
And there you go:
>>> @trace
... def div(x, y):
... return x / y
>>> div(4, 2)
enter div
leave div
2.0
We get some basic information whenever it's invoked, even when it crashes with an error:
>>> div(1, 0)
enter div
leave div
ZeroDivisionError: division by zero
All we had to do is stick a decoration on top — which is a really simple interface to affect a pretty profound change. And we can improve it further; take a look at this function:
>>> def format_call(f, *args, **kwargs):
... output = [f.__name__, '(']
... for arg in args:
... output.extend([repr(arg), ', '])
... for name, value in kwargs.items():
... output.extend([name, '=', repr(value), ', '])
... if output[-1] == ', ':
... output.pop() # remove trailing comma
... output.append(')')
... return ''.join(output)
What it does, given some function and some arguments, is organize everything in a nicely formatted string of how such an invocation might look. Now, we can have our trace decorator print a more substantial message — even print the result when the original function returns successfully, and the error when it doesn't:
>>> def trace(f):
... def wrapper(*args, **kwargs):
... call = format_call(f, *args, **kwargs)
... print('enter', call)
... try:
... result = f(*args, **kwargs)
... print('leave', call + ':', result)
... return result
... except Exception as error:
... print('leave', call + ' on error:', error)
... raise
... return wrapper
At this point, it's becoming increasingly useful — stick it on top any misbehaving function and behold the wealth of debugging information, which can really help figure out what exactly is going on:
>>> @trace
... def div(x, y):
... return x / y
>>> div(4, 2)
enter div(4, 2)
leave div(4, 2): 2.0
2.0
>>> div(1, 0)
enter div(1, 0)
leave div(1, 0) on error: division by zero
ZeroDivisionError: division by zero
Example: Memoization
Or, let's do an optimization decorator. Consider this implementation of Fibonacci's function:
>>> def fib(n):
... return n if n <= 1 else fib(n - 1) + fib(n - 2)
It's an elegant one-liner, but it's wildly inefficient. To calculate the 35th Fibonacci number, for example, we need to calculate the 34th and 33rd, right? And to calculate the 34th, we need to calculate the 33rd and 32nd. So we end up with recalculating the 33rd number twice, the 32nd four times, and the first one over a billion times, which takes almost three seconds:
>>> fib(35)
9227465 # ~3 seconds
But notice how all these reclaculations are completely unnecessary — once we figure out a number, we can cache it: store it in something like a dictionary, and simply retrieve it instantaneously whenever we need it again. This technique is called memoization, and to apply it to a function in a general way, we can write a decorator that would cache the return value by the arguments — so if the same arguments are provided, we already know the result, and don't need to invoke the function again:
>>> def cache(f):
... cache = {}
... def wrapper(*args, **kwargs):
... token = args + tuple(kwargs.items())
... if token not in cache:
... cache[token] = f(*args, **kwargs)
... return cache[token]
... return wrapper
Now, it so happens that we can't index a dictionary using a mutable object, like the keywords dictionary, for reasons we'll discuss in a later chapter; that's why instead of just using (args, kwargs) for a token, we convert the dictionary to an immutable equivalent — namely, a tuple of its items. Concatenate it with the positional arguments, and we end up with a pretty simple setup, where we only invoke the function if the token is not in the cache yet, and otherwise return the cached result. Stick it on top of our Fibonacci function:
>>> @cache
... def fib(n):
... return n if n <= 1 else fib(n - 1) + fib(n - 2)
And suddenly, it arrives at the same result in a split second, supporting numbers we couldn't hope to compute before:
>>> fib(35)
9227465 # instant
Just to play with it a little more, we can move the cache dictionary around — instead of storing it as a local variable in the decorator's scope, we can make it a function attribute in the returned wrapper:
>>> def cache(f):
... def wrapper(*args, **kwargs):
... token = args + tuple(kwargs.items())
... if token not in cache:
... wrapper.cache[token] = f(*args, **kwargs)
... return wrapper.cache[token]
... wrapper.cache = {}
... return cached
This way, after decorating fib, we can actually take a look at its cache:
>>> @cache
... def fib(n):
... return n if n <= 1 else fib(n - 1) + fib(n - 2)
>>> fib.cache
{}
Call it on something, and take a look again:
>>> fib(10)
55
>>> fib.cache
{(1,): 1, (0,): 0, (2,): 1, (3,): 2, (4,): 3, (5,): 5, (6,): 8, (7,): 13, (8,): 21, (9,): 34, (10,): 55}
We can even clear it, which will effectively force the function to recalculate the values — something that might be useful in cases where the cache might get stale and we'd like to refresh it:
>>> fib.cache.clear()
>>> fib.cache
{}
Wrapping Issues
But here we've actually stumbled unto an interesting problem. You see, that cache attribute was assigned to the wrapper, right? so the object we call fib is actually wrapper, not our original function — that one is hidden in its belly, bound by scope, but otherwise inaccessible. It sounds immaterial, but it matters; for example, if we have a nice function with an indicative name, a descriptive docstring and some type annotations, then decorating it with something (let's revert to double again) replaces it with a pretty uninspiring object, and all that valuable information is lost:
>>> @double
... def add(x: int, y: int) -> int:
... 'Returns the sum of x and y.'
... return x + y
>>> add.__name__
'wrapper'
>>> add.__doc__
None
>>> add.__annotations__
{}
To remedy this, we need to import the functools module, and use the wraps decorator on our wrapper, passing in the original function:
>>> from functools import wraps
>>> def double(f):
... @wraps(f)
... def wrapper(*args, **kwargs):
... return 2 * f(*args, **kwargs)
... return wrapper
If this is the first time you see something like this, it can be a bit much: we have a decorator, with a decoration inside it, and a function call inside that, which is just excessively meta. We'll take it apart in a moment; for now, just notice how it fixes the problem:
>>> @double
... def add(x: int, y: int) -> int:
... 'Returns the sum of x and y.'
... return x + y
>>> add.__name__
'add'
>>> add.__doc__
'returns the sum of x and y'
>>> add.__annotations__
{'x': <class 'int'>, 'y': <class 'int'>, 'return': <class 'int'>}
>>> add(1, 2)
6 # and it's still doubled!
Suddenly, we get back our old name, docstring and annotations; we even get a __wrapped__ attribute which points to the original function:
>>> add.__wrapped__
<function add at 0x...>
So yeah, decorators can be powerful — but doing them right has some nuance to it; we need to wrap things responsibly, if you know what I mean.
Second-Order Decorators
Now, to understand what exactly is happening with that weird @wraps(f) notation, we need to discuss second-order decorators. All the decorators we've seen so far were first-order; but what does this order signify, and what would it mean to bump it up?
Let's start with a simple example again. Suppose we wanted to make our double decorator more general-purpose, so we can easily get similar functionality like triple or quadruple. We might come up with a decorator factory, like so:
>>> def multiply(m):
... def decorator(f):
... def wrapper(*args, **kwargs):
... return m * f(*args, **kwargs)
... return wrapper
... return decorator
It takes some multiplier and returns a decorator, which then takes a function and returns a wrapper. When this wrapper is called, it uses both the multiplier from its factory scope and the function from its decorator scope, so everything works as expected — we can recreate double by calling multiply(2), then use it to decorate inc, and end up with the same (1 + 1) * 2 = 4 as before:
>>> double = multiply(2)
>>> @double
... def inc(x):
... return x + 1
>>> inc(1)
4
But also, we can create triple by calling multiply(3), apply that as a decoration to inc, and end up with (1 + 1) * 3 = 6:
>>> triple = multiply(3)
>>> @triple
... def inc(x):
... return x + 1
>>> inc(1)
6
What's more, we don't have to do this in two separate steps; we can just decorate inc directly with multiply(2) to double it, or directly with multiply(3) to triple it:
>>> @multiply(2)
... def inc(x):
... return x + 1
>>> inc(1)
4
>>> @multiply(3)
... def inc(x):
... return x + 1
>>> inc(1)
6
In other words, the decoration line doesn't have to be the name of a decorator — it can be any expression at all. Of course, this expression should evaluate into a decorator, because it's going to be applied to the function below it in a moment; but whatever we put there, Python will diligently apply, even if it doesn't make sense and causes an error:
>>> @1
... def f():
... pass
TypeError: 'int' object is not callable # for attempting to do 1(f)
So, while first-order decorators take a function and return a decorated version of it, second-order decorators are essentially decorator factories: they take some parameters, and return a parametrized decorator, which then gets applied to a function.
Example: Parametrized Tracing
To illustrate the point, let's parametrize our tracing decorator from before, so we can pass in a specific log object. When we want something printed to the screen, we'll pass in print; but if we have something else in mind, like writing to a file, we can pass in a different writer, and the trace decorator will use it instead.
Making this work is really not as hard as it sounds, once we get comfortable with all the layers; we just shift the previous decorator inwards, rewrite it to use the new parameters, and return its parametrized version at the end:
>>> def trace(log):
... def decorator(f):
... def wrapper(*args, **kwargs):
... call = format_call(f, *args, **kwargs)
... log('enter', call)
... try:
... result = f(*args, **kwargs)
... log('leave', call + ':', result)
... return result
... except Exception as error:
... log('leave', call + ' on error:', error)
... raise
... return wrapper
... return decorator
And it works — if we decorate with trace(print), we get the same behavior as before:
>>> @trace(print)
... def div(x, y):
... return x / y
>>> div(4, 2)
enter div(4, 2)
leave div(4, 2): 2.0
2.0
And if we decorate it with something like trace(write), where:
>>> def write(*args):
... line = ' '.join(str(arg) for arg in args) + '\n'
... open('/tmp/log.txt', 'a').write(line)
>>> @trace(write)
... def div(x, y):
... return x / y
>>> div(4, 2)
2.0
It seems like nothing happens — until we check our log file, that is:
>>> print(open('/tmp/log.txt').read())
enter div(4, 2)
leave div(4, 2): 2.0
To make this decorator production-grade, all that's left is make sure our wrapper resembles the original function, using wraps — and hopefully, at this point, it looks a little less crazy:
>>> def trace(log):
... def decorator(f):
... @wraps(f)
... def wrapper(*args, **kwargs):
... # same as before
Implementing functools.wraps
In fact, let's go ahead and implement it ourselves; it's the best way to make sure we understand something. It takes the original function, right? And it needs to return a decorator, remember — one that will be applied to the wrapper in order to fix it, to make it look like the original function:
>>> def wraps_(original):
... def fix_wrapper(wrapper):
... ?
... return fix_wrapper
Actually, the fix_wrapper decorator doesn't even need to define a wrapper of its own; we can take the wrapper that is passed in and just reset its attributes, changing its name, docstrings, annotations and everything else to be the same as the original, and adding the __wrapped__ attribute to boot:
>>> def wraps_(original):
... def fix_wrapper(wrapper):
... wrapper.__name__ = original.__name__
... wrapper.__doc__ = original.__doc__
... wrapper.__annotations__ = original.__annotations__
... wrapper.__wrapped__ = original
... return wrapper
... return fix_wrapper
At this point, using our version should be pretty idential to using the standard library, at least for those attributes:
>>> def double(f):
... @wraps_(f)
... def wrapper(*args, **kwargs):
... return 2 * f(*args, **kwargs)
... return wrapper
>>> @double
... def add(x: int, y: int) -> int:
... 'Returns the sum of x and y.'
... return x + y
>>> add.__name__
'add'
>>> add.__doc__
'returns the sum of x and y'
>>> add.__annotations__
{'x': <class 'int'>, 'y': <class 'int'>, 'return': <class 'int'>}
>>> add(1, 2)
6 # still doubled
Still More Wrapping Issues
An interesting caveat, though — and this is the case for the standard functools.wraps as well — is that no matter what attributes we fix, the signature itself is invariably broken. At the end of the day, wrapper is a function object with an (*args, **kwargs) signature, not the original one; so if we decide to inspect it, we immediately see the discrepancy:
>>> import inspect
>>> inspect.getfullargspec(add)
FullArgSpec(args=[], varargs='args', varkw='kwargs', defaults=None, kwonlyargs=[], kwonlydefaults=None, annotations={})
To have the same signature, we'd have to inspect the original and create wrapper dynamically, on the fly, using some pretty heavy metaprogramming; so before we do something like this, it's worth considering whether it's even worth it.
After all, even if we're working on function internals, where signatures really matter — we can always tell that something was wrapped if it has the __wrapped__ property; and if it does, we can peel off all the wrappers, retrieve the original function, and get the real signature from there:
>>> cursor = add
>>> while hasattr(wrapper, '__wrapped__'):
... cursor = cursor.__wrapped__
... inspect.getfullargspec(cursor)
FullArgSpec(args=['x', 'y'], varargs=None, varkw=None, defaults=None, kwonlyargs=[], kwonlydefaults=None, annotations={'return': <class 'int'>, 'x': <class 'int'>, 'y': <class 'int'>})
Incidentally, this is what inspect.signature does, which is why I had to use getfullargspec to prove my point. If we just use this new interface:
>>> inspect.signature(add)
<Signature (x: int, y: int) -> int>
It does the right thing and extracts the original parameters, not the pass-through wrapper ones.
Semi-Parametrized Decorators
At this point, if your brain hasn't melted yet, I'd like to take another poke at it, just for fun. Remember that parametrized trace decorator from before? It's a cute idea, but if we're honest, in most cases we'd probably want to just use print. However, we can't just set it as a default argument; I mean, we can — but something weird happens:
>>> def trace(log=print):
... # same as before
>>> @trace
... def inc(x):
... return x + 1
>>> inc(1)
<function wrapper() as 0x...>
Calling it returns a... wrapper?
If you think about it for a moment, it actually makes sense — trace is a second-order decorator, so when we applied it to inc, we actually passed inc as the log function, and got back a decorator! If we now apply this decorator to something else, like 1, we get a very odd wrapper, which will try to, um, increment log messages or call a number — and of course, crash.
What we forgot is a couple of parentheses at the decoration: to actually invoke trace, even with no arguments, so log defaults to print. This would indeed fix our flow, returning a decorator that is then applied to inc and replaces it with its wrapper:
>>> @trace()
... def inc(x):
... return x + 1
>>> inc(1)
enter inc(1)
leave inc(1): 2
2
Still, it's an easy mistake to make, and can be quite annoying to debug. It would be cool if we could support both cases more intuitively — like, no parentheses means default behavior, and parentheses mean parametrization. Well, take a look at this:
>>> def trace(f=None, log=print):
... if fis None:
... return lambda f: trace(f, log=log)
... def wrapper(*args, **kwargs):
... call = format_call(f, *args, **kwargs)
... log('enter', call)
... try:
... result = f(*args, **kwargs)
... log('leave', call + ':', result)
... return result
... except Exception as error:
... log('leave', call + ' on error:', error)
... raise
... return wrapper
The wrapper is almost the same as before — we even removed one level of indentation, so this would be a first-level decorator if not for the strange if-statement at its top.
To understand what it does, let's assume at first that a function is passed, so the strange clause doesn't happen. This means trace is being applied directly, with no parentheses, as a decorator:
>>> @trace
... def inc(x):
... return x + 1
Same as inc = trace(inc); so in this case, trace really does behave like a first-order decorator, simply replacing it with its traced version:
>>> inc(1)
enter inc(1)
leave inc(1): 2
2
If we do want to customize the log, however, we'll invariably have parentheses in the decoration line:
>>> @trace(log=write)
... def inc(x):
... return x + 1
In which case, the if-statement happens, and returns a temporary, intermediate lambda that expects a function. This lambda is then applied as a decorator, intercepting the function and actually calling trace again, this time with both arguments set.
This means we successfully generate and swap in the wrapper once again — while also having log set to a specific value like write, instead of the default print:
>>> inc(1)
2.0
>>> print(open('/tmp/log.txt').read())
enter inc(1)
leave inc(1): 2
To make things even clearer, we can make the function parameter position-only, since it should only be applied implicitly, through decoration (which happens by position); and the log parameter should be keyword-only, since it only gets applied explicitly, and we'd have to specify its name anyway, if only to skip over the first argument:
>>> def trace(function=None, /, *, log=print):
... # same as before
Crazy though this approach may seem, this signature is actually quite similar to the way we supported parameters in the original pattern:
>>> def compress(proceed, algorithm=None):
... # same as before
The difference was that in compress, we called the decorator directly, as in compress(save), and could easily add some arguments, like compress(save, algorithm='gzip'). Here, we can similarly do inc = trace(inc, log=write); but if we want to do it with a decoration line, we have to break it into two parts, as in inc = trace(log=write)(inc). So another way to think about what we're doing is currying, remember? Constructing a function gradually, making trace(log=write) on its own just a partial step.
The functools module where we found wraps can actually provide some help with that, too. We can take a function and some arguments, and fix those values in its suspended call, creating a similar and appropriately named partial object:
>>> from functools import partial
>>> f = partial(add, 1) # fixes add with x = 1
>>> f(2) # completes with y = 2
3
We can use this instead of our improvised lambda, which may be somewhat clearer:
>>> def trace(function=None, /, *, log=print):
... if function is None:
... return partial(trace, log=log)
... # same as before
And that should be enough decorators for one day.
Decorators: Conclusion
We started with a design pattern, which was intended to let us compose behaviors dynamically, assuming they have the same, specific signature. I think it helps explain how this way of thinking came about; but as we saw, Python decorators can and do much more. Since we have generic, catch-all signatures, we can use decorators to address cross-cutting concerns, effectively equipping any function with superpowers using a single, simple statement on top.
This statement can be a first-order decorator, which swaps the function for an enhanced wrapper; or, it can be an expression that calls a second-order decorator to further parametrize the decoration; and we even saw a semi-parametrized decorator that goes both ways. We talked about how to wrap things responsibly, and explored examples like tracing and memoization.
There's obviously much more to tell about this fascinating approach — we'll do some of it in the exercises and some of it in future chapters on software design, where we'll often encounter decorators in one way or another.