Search code examples
pythonpython-3.xsignaturefunctoolspydoc

Overriding function signature (in help) when using functools.wraps


I'm creating a wrapper for a function with functools.wraps. My wrapper has the effect of overriding a default parameter (and it doesn't do anything else):

def add(*, a=1, b=2):
   "Add numbers"
   return a + b

@functools.wraps(add)
def my_add(**kwargs):
    kwargs.setdefault('b', 3)
    return add(**kwargs)

This my_add definition behaves the same as

@functools.wraps(add)
def my_add(*, a=1, b=3):
    return add(a=a, b=b)

except that I didn't have to manually type out the parameter list.

However, when I run help(my_add), I see the help string for add, which has the wrong function name and the wrong default argument for the parameter b:

add(*, a=1, b=2)
    Add numbers

How can I override the function name and the default argument in this help() output?

(Or, is there a different way to define my_add, using for example some magic function my_add = magic(add, func_name='my_add', kwarg_defaults={'b': 3}) that will do what I want?)


Solution

  • Let me try and explain what happens.

    When you call the help functions, this is going to request information about your function using the inspect module. Therefore you have to change the function signature, in order to change the default argument.

    Now this is not something that is advised, or often preferred, but who cares about that right? The provided solution is considered hacky and probably won't work for all versions of Python. Therefore you might want to reconsider how important the help function is... Any way let's start with some explanation on how it was done, followed by the code and test case.


    Copying functions

    Now the first thing we will do is copy the entire function, this is because I only want to change the signature of the new function and not the original function. This decouples the new my_add signature (and default values) from the original add function.

    See:

    For ideas of how to do this (I will show my version in a bit).

    Copying / updating signature

    The next step is to get a copy of the function signature, for that this post was very useful. Except for the part where we have to adjust the signature parameters to match the new keyword default arguments.

    For that we have to change the value of a mappingproxy, which we can see when running the debugger on the return value of inspect.signature(g). Now so far this can only be done by changing the private variables (the values with leading underscores _private). Therefore this solution will be considered hacky and is not guaranteed to withstand possible updates. That said, let's see the solution!


    Full code

    import inspect
    import types
    import functools
    
    
    def update_func(f, func_name='', update_kwargs: dict = None):
        """Based on http://stackoverflow.com/a/6528148/190597 (Glenn Maynard)"""
        g = types.FunctionType(
                code=f.__code__,
                globals=f.__globals__.copy(),
                name=f.__name__,
                argdefs=f.__defaults__,
                closure=f.__closure__
        )
    
        g = functools.update_wrapper(g, f)
        g.__signature__ = inspect.signature(g)
        g.__kwdefaults__ = f.__kwdefaults__.copy()
    
        # Adjust your arguments
        for key, value in (update_kwargs or {}).items():
            g.__kwdefaults__[key] = value
            g.__signature__.parameters[key]._default = value
    
        g.__name__ = func_name or g.__name__
        return g
    
    
    def add(*, a=1, b=2):
        "Add numbers"
        return a + b
    
    my_add = update_func(add, func_name="my_add", update_kwargs=dict(b=3))
    

    Example

    if __name__ == '__main__':
        a = 2
        
    
        print("*" * 50, f"\nMy add\n", )
        help(my_add)
    
        print("*" * 50, f"\nOriginal add\n", )
        help(add)
    
        print("*" * 50, f"\nResults:"
                        f"\n\tMy add      : a = {a}, return = {my_add(a=a)}"
                        f"\n\tOriginal add: a = {a}, return = {add(a=a)}")
    

    Output

    ************************************************** 
    My add
    
    Help on function my_add in module __main__:
    
    my_add(*, a=1, b=3)
        Add numbers
    
    ************************************************** 
    Original add
    
    Help on function add in module __main__:
    
    add(*, a=1, b=2)
        Add numbers
    
    ************************************************** 
    Results:
        My add      : a = 2, return = 5
        Original add: a = 2, return = 4
    

    Usages

    • f: is the function that you want to update
    • func_name: is optionally the new name of the function (if empty, keeps the old name)
    • update_kwargs: is a dictionary containing the key and value of the default arguments that you want to update.

    Notes

    • The solution is using copy variables to make full copies of dictionaries, such that there is no impact on the original add function.
    • The _default value is a private variable, and can be changed in future releases of python.