Signature-changing decorator: properly documenting additional argument

Question

Let's say I have a custom decorator, and I want it to properly handle docstring of decorated function. Problem is: my decorator adds an argument.

from functools import wraps

def custom_decorator(f):
    @wraps(f)
    def wrapper(arg, need_to_do_more):
        '''
        :param need_to_do_more: if True: do more
        '''
        args = do_something(arg)

        if need_to_do_more:
            args = do_more(args)

        return f(args)

    return wrapper

You can see the argument is not actually passed to decorated function, but used by the wrapper - which may or may not be relevant here.

How can I properly handle documenting the additional argument? Is it a good practice for a wrapper to take an additional argument, or should I avoid it?

Or should I rather use a different solution, like:

• making wrapper a simple higher order function, with the function it calls passed as the third argument
• refactoring the wrapper into two separate functions?

Answer 1

So - __doc__ apart this is tricky - and, due to more and more developers relying on the automatic parameter suggestions when coding, which is given by IDE introspection, it is really needed for any decorator that will add extra named parameters to a function.

I got to that in a project I am developing, and the solution is to create a new, dummy function, that will have the desired combined signature to be shown - and then using this new dummy function as the parameter to the @wraps call,.

Here is my code - it is good enough, and so unrelated with the other project I will likely put it in a decorators Python package soon. For now:

def combine_signatures(func, wrapper=None):
    """Adds keyword-only parameters from wrapper to signature
    
    Use this in place of `functools.wraps` 
    It works by creating a dummy function with the attrs of func, but with
    extra, KEYWORD_ONLY parameters from 'wrapper'.
    To be used in decorators that add new keyword parameters as
    the "__wrapped__"
    
    Usage:
    
    def decorator(func):
        @combine_signatures(func)
        def wrapper(*args, new_parameter=None, **kwargs):
            ...
            return func(*args, **kwargs)
    """
    # TODO: move this into 'extradeco' independent package
    from functools import partial, wraps
    from inspect import signature, _empty as insp_empty, _ParameterKind as ParKind
    from itertools import groupby

    if wrapper is None:
        return partial(combine_signatures, func)

    sig_func = signature(func)
    sig_wrapper = signature(wrapper)
    pars_func = {group:list(params)  for group, params in groupby(sig_func.parameters.values(), key=lambda p: p.kind)}
    pars_wrapper = {group:list(params)  for group, params in groupby(sig_wrapper.parameters.values(), key=lambda p: p.kind)}

    def render_annotation(p):
        return f"{':' + (repr(p.annotation) if not isinstance(p.annotation, type) else repr(p.annotation.__name__)) if p.annotation != insp_empty else ''}"

    def render_params(p):
        return f"{'=' + repr(p.default) if p.default != insp_empty else ''}"

    def render_by_kind(groups, key):
        parameters = groups.get(key, [])
        return [f"{p.name}{render_annotation(p)}{render_params(p)}" for p in parameters]

    pos_only = render_by_kind(pars_func, ParKind.POSITIONAL_ONLY)
    pos_or_keyword = render_by_kind(pars_func, ParKind.POSITIONAL_OR_KEYWORD)
    var_positional = [p for p in pars_func.get(ParKind.VAR_POSITIONAL,[])]
    keyword_only = render_by_kind(pars_func, ParKind.KEYWORD_ONLY)
    var_keyword = [p for p in pars_func.get(ParKind.VAR_KEYWORD,[])]

    extra_parameters = render_by_kind(pars_wrapper, ParKind.KEYWORD_ONLY)

    def opt(seq, value=None):
        return ([value] if value else [', '.join(seq)]) if seq else []

    annotations = func.__annotations__.copy()
    for parameter in pars_wrapper.get(ParKind.KEYWORD_ONLY):
        annotations[parameter.name] = parameter.annotation

    param_spec = ', '.join([
        *opt(pos_only),
        *opt(pos_only, '/'),
        *opt(pos_or_keyword),
        *opt(keyword_only or extra_parameters, ('*' if not var_positional else f"*{var_positional[0].name}")),
        *opt(keyword_only),
        *opt(extra_parameters),
        *opt(var_keyword, f"**{var_keyword[0].name}" if var_keyword else "")
    ])
    declaration = f"def {func.__name__}({param_spec}): pass"

    f_globals = func.__globals__
    f_locals = {}

    exec(declaration, f_globals, f_locals)

    result = f_locals[func.__name__]
    result.__qualname__ = func.__qualname__
    result.__doc__ = func.__doc__
    result.__annotations__ = annotations

    return wraps(result)(wrapper)

Testing in interactive mode one gets this result:

IPython 7.8.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: from terminedia.utils import combine_signatures                                                                                                    

In [2]: def add_color(func): 
   ...:     @combine_signatures(func) 
   ...:     def wrapper(*args, color=None, **kwargs): 
   ...:         global context 
   ...:         context.color = color 
   ...:         return func(*args, **kw) 
   ...:     return wrapper 
   ...:                                                                                                                                                    

In [3]: @add_color 
   ...: def line(p1, p2): 
   ...:     pass 
   ...:                                                                                                                                                    

In [4]: line                                                                                                                                               
Out[4]: <function __main__.line(p1, p2, *, color=None)>

(As for doc strings, as in the question - once one have got all the wrapper and function data, it is a matter of text handling before pasting result.__doc__ = func.__doc__ . Since each project will have different styles for documenting parameters inside docstrings, it can't be done reliably in a 'one size fits all', but with some string splicing and testing it could be perfected for any given doc string style)

Answer 2

If the result of the decoration is expected to always supply the same thing to that argument, I'd recommend making it a parameterized decorator. I'm guessing you thought of that, but it needed to be said.

Other than that, I would definitely recommend breaking it in two, like your second suggestion. Then the user of the decorators can provide "overloaded" versions (not truly overloaded, since they will need different names) that use the two different decorators.

Another possible option is giving the parameter a default value.

Lastly, if you simply must keep it the way that it is, you need to append the new parameter documentation to the end of __doc__ after the wrapper definition.

So your example (shortened) would look like this:

def custom_decorator(f):
    @wraps(f)
    def wrapper(arg, need_to_do_more):
        ...

    wrapper.__doc__ += "/n:param need_to_do_more: if True: do more"
    return wrapper

This is because the @wraps(f) decorator replaces wrapper's documentation with f's. Adding it afterwards will actually combine the two.

The other option along these lines is to document custom_decorator so that it says that wrapped methods need to add the parameter to their documentation. This, as well as the splitting up of the decorator, put the burden on the user, but they make the intention more explicit ("...Explicit is better than implicit..." - Zen of Python)