Python decorators#
Python decorators are used to add functionality to functions, methods or classes. Some of the use cases are e.g.:
- add logging/tracing around function calls
- add locking for threaded code
- register functions or classes
- add caching to callables
Decorating functions#
A very basic decorator to trace function entry and exit could look like this:
>>> def trace(func):
... # create a wrapper function...
... def _wrapper(*args, **kwargs):
... # ...that prints entry to and exit from the wrapped function,
... # with function name, arguments and result...
... print(f'--> {func.__name__}(args={args}, kwargs={kwargs})')
... result = func(*args, **kwargs)
... print(f'<-- {func.__name__} -> {result}')
... # ...and return the wrapper for use instead of the original function
... return _wrapper
...
We can now apply this decorator to a function:
Calling our decorated "traced" function:
Note how we have instrumented the original function with tracing output, while not modifying any of the original function code.
Since a (function) decorator is just a callable that takes a function as an
argument and returns an enhanced, wrapped version of that function, the
@
-decorator syntax is merely syntactial "sugar" for:
>>> def inc(x):
... """Return x increased by 1."""
... return x + 1
...
>>> inc = trace(inc)
>>> inc(6)
--> inc(args=(6,), kwargs={})
<-- inc -> 7
>>>
Let's take a look at how our decorated function behaves, now using @decorator-syntax again:
>>> def trace(func):
... # create a wrapper function...
... def _wrapper(*args, **kwargs):
... # ...that prints entry to and exit from the wrapped function,
... # with function name, arguments and result...
... print(f'--> {func.__name__}(args={args}, kwargs={kwargs})')
... result = func(*args, **kwargs)
... print(f'<-- {func.__name__} -> {result}')
... # ...and return the wrapper for use instead of the original function
... return _wrapper
...
>>> @trace
... def inc(x):
... """Return x increased by 1.
... """
... return 1
...
>>> print(inc.__doc__)
None
>>> import inspect
>>> inspect.signature(inc)
<Signature (*args, **kwargs)>
>>>
Hm, this doesn't look too good a citizen:
- the original function documentation has been lost
- there's no information about the original function argument signature
This is due to the fact that we see docstring and function signature of the
wrapper, not the wrapped function. Luckily, there's a convenient way for us to
retain this information using the functools
library:
>>> import functools, inspect
>>> def trace(func):
... # create a wrapper function...
... @functools.wraps(func)
... def _wrapper(*args, **kwargs):
... # ...that prints entry to and exit from the wrapped function,
... # with function name, arguments and result...
... print(f'--> {func.__name__}(args={args}, kwargs={kwargs})')
... result = func(*args, **kwargs)
... print(f'<-- {func.__name__} -> {result}')
... # ...and return the wrapper for use instead of the original function
... return _wrapper
...
>>> @trace
... def inc(x):
... """Return x increased by 1.
... """
... return x + 1
...
>>> print(inc.__doc__)
Return x increased by 1.
>>> inspect.signature(inc)
<Signature (x)>
>>>
Decorating methods#
We can just as well decorate methods:
>>> class Increaser:
... def __init__(self, increment=1):
... self.increment = increment
... @trace
... def inc(self, x):
... """Return x increased by increment init argument.
... """
... return x + self.increment
...
>>> inc = Increaser(3)
>>> inc.inc(10)
--> inc(args=(<__main__.Increaser object at 0x7fc62185c320>, 10), kwargs={})
<-- inc -> 13
>>>
>>> print(inc.inc.__doc__)
Return x increased by increment init argument.
>>> inspect.signature(inc.inc)
<Signature (x)>
>>>
In fact, Python implements its notion of class methods and static methods
via the built-in @classmethod
and
@staticmethod
decorators.
Decorating classes#
Decorating classes is just as easy:
>>> class Registry:
... _registered_classes = []
...
... @classmethod
... def register(cls, register_cls):
... cls._registered_classes.append(register_cls)
...
... @classmethod
... def registered_classes(cls):
... return cls._registered_classes
...
>>> register = Registry.register
>>> registry = Registry()
>>>
>>> @register
... class MyClass1:
... pass
...
>>> @register
... class MyClass2:
... pass
...
>>> registry.registered_classes()
[<class '__main__.MyClass1'>, <class '__main__.MyClass2'>]
>>>
Advanced Decorators#
Decorators can also take arguments.
Let's modify the tracing functionality to allow for selective entry and/or exit tracing:
>>> import functools, inspect
>>> def traced(entry=True, exit=True):
... # create tracing wrappers depending on the entry & exit args
... if entry and exit:
... def trace(func):
... @functools.wraps(func)
... def _wrapper(*args, **kwargs):
... print(f'--> {func.__name__}(args={args}, kwargs={kwargs})')
... result = func(*args, **kwargs)
... print(f'<-- {func.__name__} -> {result}')
... return _wrapper
...
... elif entry:
... def trace(func):
... @functools.wraps(func)
... def _wrapper(*args, **kwargs):
... print(f'--> {func.__name__}(args={args}, kwargs={kwargs})')
... result = func(*args, **kwargs)
... return _wrapper
...
... elif exit:
... def trace(func):
... @functools.wraps(func)
... def _wrapper(*args, **kwargs):
... result = func(*args, **kwargs)
... print(f'<-- {func.__name__} -> {result}')
... return _wrapper
... else:
... trace = None
...
... # create the decorator that will add the selected tracing setup
... def decorate(cls):
... if trace is not None:
... for (name, method) in inspect.getmembers(cls):
... if not name.startswith('__'):
... wrapped = trace(method)
... setattr(cls, name, wrapped)
... return cls
...
... return decorate
...
>>>
We'll try this out on a delicious fruit salad:
>>> class FruitSalad:
...
... def __init__(self):
... self.fruits = {}
...
... def add(self, fruit, weight):
... if fruit not in self.fruits:
... self.fruits[fruit] = weight
... else:
... self.fruits[fruit] += weight
...
>>>
>>> fruit_salad = FruitSalad()
>>> fruit_salad.add('apple', '500')
>>> fruit_salad.add('orange', '800')
>>>
Make a traced fruit salad:
>>>
>>> @traced()
... class FruitSalad:
...
... def __init__(self):
... self.fruits = {}
...
... def add(self, fruit, weight):
... if fruit not in self.fruits:
... self.fruits[fruit] = weight
... else:
... self.fruits[fruit] += weight
...
>>> fruit_salad = FruitSalad()
>>> fruit_salad.add('apple', '500')
--> add(args=(<__main__.FruitSalad object at 0x7fc61868cda0>, 'apple', '500'),
kwargs={})
<-- add -> None
>>> fruit_salad.add('orange', '800')
--> add(args=(<__main__.FruitSalad object at 0x7fc61868cda0>, 'orange', '800'),
kwargs={})
<-- add -> None
>>>
This might look a little bit confusing initially. The important point is to
realize that the callable traced
we use in the @traced()
decoration line is
now rather a "decorator factory" than a decorator: it creates a decorator
(the decorate
function), and returns this to be applied on the decorated
class.
We can now switch off tracing method entry and only trace funtion exit:
>>>
>>> @traced(entry=False)
... class FruitSalad:
...
... def __init__(self):
... self.fruits = {}
...
... def add(self, fruit, weight):
... if fruit not in self.fruits:
... self.fruits[fruit] = weight
... else:
... self.fruits[fruit] += weight
...
>>> fruit_salad = FruitSalad()
>>> fruit_salad.add('apple', '500')
<-- add -> None
>>> fruit_salad.add('orange', '800')
<-- add -> None
>>>
Note: Maybe that's just cosmetics but it's a little annoying that we need
to use @traced()
syntax even if we simple use the default arguments.
Let's make use of Python's [keyword-only] syntax to work around this:
>>> def traced(cls=None, *, entry=True, exit=True):
... if entry and exit:
... def trace(func):
... @functools.wraps(func)
... def _wrapper(*args, **kwargs):
... print(f'--> {func.__name__}(args={args}, kwargs={kwargs})')
... result = func(*args, **kwargs)
... print(f'<-- {func.__name__} -> {result}')
... return _wrapper
...
... elif entry:
... def trace(func):
... @functools.wraps(func)
... def _wrapper(*args, **kwargs):
... print(f'--> {func.__name__}(args={args}, kwargs={kwargs})')
... result = func(*args, **kwargs)
... return _wrapper
...
... elif exit:
... def trace(func):
... @functools.wraps(func)
... def _wrapper(*args, **kwargs):
... result = func(*args, **kwargs)
... print(f'<-- {func.__name__} -> {result}')
... return _wrapper
... else:
... trace = None
...
... def decorate(cls):
... if trace is not None:
... for (name, method) in inspect.getmembers(cls):
... if not name.startswith('__'):
... wrapped = trace(method)
... setattr(cls, name, wrapped)
... return cls
...
...
... if cls is None:
... # called with arguments
... return decorate
... else:
... # invoked without arguments
... return decorate(cls)
...
>>>
We are now able to omit () from the decoration line:
>>> @traced
... class FruitSalad:
...
... def __init__(self):
... self.fruits = {}
...
... def add(self, fruit, weight):
... if fruit not in self.fruits:
... self.fruits[fruit] = weight
... else:
... self.fruits[fruit] += weight
...
>>> fruit_salad = FruitSalad()
>>> fruit_salad.add('apple', '500')
--> add(args=(<__main__.FruitSalad object at 0x7fc61868cda0>, 'apple', '500'),
kwargs={})
<-- add -> None
>>> fruit_salad.add('orange', '800')
--> add(args=(<__main__.FruitSalad object at 0x7fc61868cda0>, 'orange', '800'),
kwargs={})
<-- add -> None
>>>
Further Reading#
One of the most exhaustive discussions of decorator intricacies is a series of blog posts by Graham Dumpleton, who also wrote the wrapt library for decoration purposes.