Python decorators - how to keep function signature (set of arguments) unchanged

Problem

Decorators modify original function.

As a result we have new function with different sugnature (set of arguments).

Sometimes we want to keep old set of arguments. For example if we use some tool that based on the function arguments it will cease to work because new function will have other set of arguments.

For example after:

def my_decorator(original_function):
    def wrapper(*args, **kwargs):
        try:
            return original_function(*args, **kwargs)
        finally:
            some_clean_up()
    return wrapper

We will have only faceless *args, **kwargs as our new function arguments.

How to save function arguments list after it had been decorated

In Python 3.3 and above (PEP0362) you can use special attribute __signature__ to save original arguments list:

import inspect


def my_decorator(original_function):
    def wrapper(*args, **kwargs):
        try:
            return original_function(*args, **kwargs)
        finally:
            some_clean_up()
    wrapper.__signature__ = inspect.signature(original_function)
    return wrapper

How to save other function attributes

The arguments list is not the only function feature that we lost after our decorator.

For example we will loose also __doc__ and that will be a problem for doc-tests - all the doc-tests won’t work anymore.

This and many other attributes of the function that had been decorated you can save with help of wraps from functools:

import inspect
import functools


def my_decorator(original_function):
    @functools.wraps
    def wrapper(*args, **kwargs):
        try:
            return original_function(*args, **kwargs)
        finally:
            some_clean_up()
    wrapper.__signature__ = inspect.signature(original_function)
    return wrapper

It’s a pity that for __signature__ you need separate code and cannot use the same wraps.

This is because many functions just do not have __signature__ attribute. You have to create it with help of inspect. wraps just copy existing attributes so it’s no help there.

Written on November 16, 2018