Nate Tracy-Amoroso

Home > Python decorators that have arguments and preserve docstrings

tl;dr

# Outermost decorator accepts arguments
def factory(the_arg): 
    # Inner decorator accepts the function
    def decorator(_func): 
        # Decorator creates this function 'inner' to replace the original
        def inner(*args, **kwargs): 
            print(f"Pre function : {the_arg}")
            return_val = _func(*args, **kwargs)
            print(f"Post function : {the_arg}")
            return return_val
        # Assign the original functions docstring to the new version 'inner'
        inner.__doc__ = _func.__doc__ 
        return inner
    return decorator

@factory("some_arg")
def new_function():
    print("Hello from new_function")

new_function()
> python3 example.py
Pre function : some_arg
Hello from new_function
Post function : some_arg

I like using doctest for python3 to write tests. It looks like this

def add(a,b):
    """
    >>> add(5,6)
    11
    """
    return a + b

Then you can use python3 -m doctests blah.py and it will run add(5,6) and compare it to 11. Great easy inline tests just like rust.

The problem is in a case like this

# Add the_arg to whatever the original function returns
def add_to_result(the_arg):
    def decorator(_func):
        def inner(*args, **kwargs):
            return _func(*args, **kwargs) + the_arg
        return inner
    return decorator

@add_to_result(10)
def add(a, b):
    """
    >>> add(1,2)
    13
    """
    return a + b

doctest won’t be able to find that there is a docstring on the add function because the decorator threw through it away here:

def inner(*args, **kwargs): 
    return _func(*args, **kwargs) + the_arg # _func has the docstring
return inner # but inner doesn't becuase it's a brand new function

All we have to do is add one line inner.__doc__ = _func.__doc__ to our decorator and we can preserve it.