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)