Python decorator show signature and docstring from both decorator and decorated function

I'm trying to write a decorator whose wrapper function takes an additional input argument, based on which some processing is done within the wrapper. However, I'm having a hard time a) figuring out how to show this additional parameter in the signature, especially with parameter previews/helper (tested on ipynb and PyCharm) b) Figuring out how to show the info from both docstrings.

The setup would be something like this

from functools import wraps

def my_decorator(func):
    @wraps(func)
    def wrapper(self, *args, fancy_processing=False, **kwargs):
        """Some nice docstring
        :param fancy_processing: If set, does fancy processing before executing the function
        """
        output = func(self, *args, **kwargs)
        if fancy_processing:
            output += 1 # some dummy action
        return output
    return wrapper


def foo(important_param, other_important_param=4):
    """This is a well-written docstring.
    :param important_param: Super important parameter
    :param other_important_param: Other super important parameter
    :returns: Something incredibly useful
    """
    print(important_param)  # dummy action
    return (other_important_param +1)  # some other dummy action

Now depending on whether I use wraps or not I get the signature and docstring of my wrapper function or my foo function, but not both. For the docstring, I could do something like

...
# @wraps(func) # don't use the `wraps` functionality
...
    return output
wrapper.__doc__ += f"\n {func.__doc__}"
return wrapper

but this doesn't seem very well-formated and also doesn't solve the issue with the signature not reflecting parameters from both functions. How would I go about solving this?


Solution 1:

If you look at what functools.wraps is eant to do, and how it is implemented, you can notice it is pretty simple : it just replaces some attributes of the wrapper by some of the wrapped, so that the wrapper looks like the wrapped.

If you want to have a fancy merge between two docstrings, either write it yourself, or if you want to use this decorator for many functions, create a function to do it automatically. Here is a proof of concept :

import docstring_parser  # installed with PIP


def merge_docstrings(decorated_func, wrapper_func) -> str:
    reST_style = docstring_parser.Style.REST
    decorated_doc = docstring_parser.parse(decorated_func.__doc__, style=reST_style)
    wrapper_doc = docstring_parser.parse(wrapper_func.__doc__, style=reST_style)
    # adding the wrapper doc into the decorated doc
    decorated_doc.short_description += "\n" + wrapper_doc.short_description
    return_index = next((index for index, meta_elem in enumerate(decorated_doc.meta)
                         if isinstance(meta_elem, docstring_parser.DocstringReturns)
                         ), 0)
    for offset, wrapper_doc_param in enumerate(wrapper_doc.params):
        decorated_doc.meta.insert(return_index + offset, wrapper_doc_param)

    return docstring_parser.compose(decorated_doc, style=reST_style)


def my_decorator(func):
    def wrapper(self, *args, fancy_processing=False, **kwargs):
        """Some nice docstring
        :param fancy_processing: If set, does fancy processing before executing the function
        """
        output = func(self, *args, **kwargs)
        if fancy_processing:
            output += 1 # some dummy action
        return output
    wrapper.__doc__ = merge_docstrings(func, wrapper)
    return wrapper


def foo(important_param, other_important_param=4):
    """This is a well-written docstring.
    :param important_param: Super important parameter
    :param other_important_param: Other super important parameter
    :returns: Something incredibly useful
    """
    print(important_param)  # dummy action
    return (other_important_param +1)  # some other dummy action


decorated_foo = my_decorator(foo)  # the alternative way to decorate a function, and allows to keep the original too


assert foo.__doc__ == ('This is a well-written docstring.\n    '
                       ':param important_param: Super important parameter\n    '
                       ':param other_important_param: Other super important parameter\n    '
                       ':returns: Something incredibly useful\n    ')
assert decorated_foo.__doc__ == \
                      ('This is a well-written docstring.\n'
                       'Some nice docstring\n'
                       ':param important_param: Super important parameter\n'
                       ':param other_important_param: Other super important parameter\n'
                       ':param fancy_processing: If set, does fancy processing before executing the function\n'
                       # ^^^ this one has been added
                       ':returns: Something incredibly useful')
help(decorated_foo)
# wrapper(self, *args, fancy_processing=False, **kwargs)
#     This is a well-written docstring.
#     Some nice docstring
#     :param important_param: Super important parameter
#     :param other_important_param: Other super important parameter
#     :param fancy_processing: If set, does fancy processing before executing the function
#     :returns: Something incredibly useful