Search code examples
pythonpython-decoratorsdocstring

Preserving function docstrings when decorating them using decorator classes


I wonder what is the correct way to write a decorator class such that it preserves the docstring of the decorated function. For example, consider the following lines of code:

import functools

class _vectorize:
    """ (private) decorator to vectorize a function """


    def __init__(self, func):
        functools.update_wrapper(self, func)
        self._func = func

    def __call__(self, values, *args, **kwargs):

        return [self._func(val, *args, **kwargs) for val in values]


@_vectorize
def multiply(v, s):
    """ Multiplies all values in `v` by scalar `s`. """
    return v*s

I would like that the client sees the docstring of multiply function by calling help(multiply) but this is not the case:

>>> help(multiply)

Help on instance of _vectorize in module docstring:

multiply = class _vectorize
 |  (private) decorator to vectorize a function
 |  
 |  Methods defined here:
 |  
 |  __call__(self, values, *args, **kwargs)
 |  
 |  __init__(self, func)

To see the correct documentation one has to directly print the __doc__ attribute:

>>> print(multiply.__doc__)
 Multiplies all values in `v` by scalar `s`. 

How should I modify the decorator so that help prints the desired documentation?


Solution

  • I use functools.wraps which is shorter and works in your case and displays proper docstring:

    import functools
    
    def _vectorize(func):
            @functools.wraps(func)
            def wrapFunc(values, *args, **kwargs):
                    return [func(val, *args, **kwargs) for val in values]
            return wrapFunc
    
    @_vectorize
    def multiply(v, s):
        """ Multiplies all values in `v` by scalar `s`. """
        return v*s
    

    See: What does functools.wraps do?