Sunday, 18 September 2011

Propagating Python function annotations through decorators

When we decorate a Python function we must take some care to ensure that the decorated function can present itself in the same way as before its decoration. Let's explore what is needed.

Decorating a Python function

Say we have an undecorated function f(). It features a docstring (line 2) and annotated arguments:

def f(a:str):
	"an undecorated function"
	return a

print(f.__name__,f.__annotations__,f.__doc__,f('foo'))

The final line in the previous snippet prints the name of the function, its annotation dictionary and the docstring. The output looks like this:

f {'a': <class 'str'>} an undecorated function foo

Decorating a Python function

Now say we have a function do_something() and we would like to have a decorator that would arrange that a function would call do_something() first. We might set this up like this:

def do_something():
	print('oink')
	
def decorator(f):
	def g(*args,**kwargs):
		do_something()
		return f(*args,**kwargs)
	return g
	
@decorator
def f(a:str):
	"a naively decorated function"
	return a

print(f.__name__,f.__annotations__,f.__doc__,f('foo'))

The final line would produce the following output:

oink
g {} None foo

It is clear that do_something() is executed as it does print oink but although our decorated function can be called as f(), its name, annototations and doctring are different.

Preserving attributes when decorating a Python function

To resolve these matters we can rewrite our decorator as follows:

def decorator(f):
	def g(*args,**kwargs):
		do_something()
		return f(*args,**kwargs)
	g.__name__ = f.__name__
	g.__doc__ = f.__doc__
	g.__annotations__ = f.__annotations__
	return g
	
@decorator
def f(a:str):
	"a less naively decorated function"
	return a

print(f.__name__,f.__annotations__,f.__doc__,f('foo'))

We now simply copy the relevant attributes to the newly created function and if we now examine the output produced by the final line we get what we expect:

oink
f {'a': <class 'str'>} a less naively decorated function foo

No comments:

Post a Comment