9
\$\begingroup\$

I ran across an issue that proved to be more complicated than I thought: Changing an arbitrary argument of a function purely based upon its name in Python - possibly via a decorator. I tried to implement a generic solution, taking into account Python's positional arguments, keyword arguments and argument defaults.

First things first, type hints purely for readability:

from typing import Any, Callable, Dict, List, Tuple

The following function is the core. It looks at func, the target function, and tries to locate the argument named arg_name. It then applies task to it.

def _change_one_argument(
    func: Callable,
    arg_name: str,
    task: Callable,
    args: Tuple,
    kwargs: Dict,
) -> Tuple[Tuple, Dict]:
    """
    Apply `task` to one argument in `args` or `kwargs` of `func`
    based on `argument`'s name.
    """
    
    arg_names = func.__code__.co_varnames[:func.__code__.co_argcount]
    arg_position = arg_names.index(arg_name)
    
    if arg_position < len(args): # from args
        old_value = args[arg_position]
        args = list(args) # for future change and as a flag
    elif arg_name in kwargs.keys(): # from kwargs
        old_value = kwargs.get(arg_name)
    else: # from default
        offset = len(arg_names) - len(func.__defaults__)
        old_value = func.__defaults__[arg_position - offset]

    new_value = task(old_value)

    if isinstance(args, list): # consider this as a flag
        args[arg_position] = new_value
        args = tuple(args) # change back to tuple
    else:
        kwargs = kwargs.copy() # beware the mutable dict
        kwargs[arg_name] = new_value

    return args, kwargs

As potential use case for testing, I built a decorator, which accepts an arbitrary number of keyword arguments. Each decorator argument represents one argument of the decorated function that is to be changed. The actual change is handled by a callable passed as the argument to the decorator.

def change_arguments(**changes: Callable):
    """
    Name parameters via keywords and pass callables,
    which are applied to `func`'s arguments.
    """
    def outer(func: Callable):
        def inner(*args: Any, **kwargs: Any):
            for arg_name, task in changes.items():
                args, kwargs = _change_one_argument(
                    func = func,
                    arg_name = arg_name,
                    task = task,
                    args = args,
                    kwargs = kwargs,
                )
            return func(*args, **kwargs)
        return inner
    return outer

One can test it for instance as follows:

@change_arguments(c = lambda x: x * 100)
def foo(a: int, b: int = 77, c: int = 88, d: int = 99):
    e = 14
    print(a, b, c, d, e)

foo(10, 11)
foo(10, 11, d = 13)
foo(10, 11, 12, 13)
foo(10, 11, c = 12, d = 13)
foo(10, 11, d = 13, c = 12)

It produces:

10 11 8800 99 14
10 11 8800 13 14
10 11 1200 13 14
10 11 1200 13 14
10 11 1200 13 14

I know that the error handling is not perfect. My solution also does not help every time if the decorated function is written in C and does, possibly, not accept keyword arguments. It also heavily relies on the fact that the argument has a name which is not, strictly speaking, always the case. Last but not least, my handling of mutable keyword argument dictionaries could be better. In the above case, it does not matter, but if _change_one_argument was a directly callable API of some sort, I am not sure - maybe it should not return anything, accept args only as a list and apply changes directly to args and kwargs (no copies made).

\$\endgroup\$

1 Answer 1

11
\$\begingroup\$

You can use inspect.signature to simplify your code.

>>> def foo(a, /, b, *, c, d=1):
...     ...
... 
>>> import inspect
>>> args = inspect.signature(foo).bind(1, 2, c=3)
>>> args.apply_defaults()
>>> args
<BoundArguments (a=1, b=2, c=3, d=1)>
>>> args.arguments["c"] = 10 * args.arguments["c"]
>>> args
<BoundArguments (a=1, b=2, c=30, d=1)>
>>> args.args, args.kwargs
((1, 2), {'c': 30, 'd': 1})

As such we can drastically reduce the amount of code you need. Effectively removing _change_one_argument. For example you can change the code to:

def change_arguments(**changes):
    def outer(func):
        def inner(*args, **kwargs):
            _args = signature.bind(*args, **kwargs)
            _args.apply_defaults()
            for name, changer in changes.items():
                _args.arguments[name] = changer(_args.arguments[name])
            return func(*_args.args, **_args.kwargs)
        signature = inspect.signature(func)
        return inner
    return outer
\$\endgroup\$
4
  • \$\begingroup\$ Cool. Knowing about inspect.signature could have essentially saved me some trouble ... Thanks a lot! \$\endgroup\$
    – s-m-e
    Dec 16, 2021 at 22:16
  • 2
    \$\begingroup\$ @s-m-e I'm happy I could help. However, you may want to unaccept my answer, as leaving the accept open indicates you're open for more (and better) reviews of your code. I've focused on one aspect of your code, and I know there are some more comments which can be made, for example your type hints could have a couple of comments. \$\endgroup\$
    – Peilonrayz
    Dec 17, 2021 at 0:11
  • \$\begingroup\$ This also works for positional arguments. \$\endgroup\$
    – RootTwo
    Dec 17, 2021 at 2:51
  • \$\begingroup\$ @RootTwo Yes. The change should work for all the kinds of arguments; positional only, positional or keyword, var positional (*args), keyword only, and var keyword (**kwargs). The OP may need to change how var positional and/or var keyword works depending on use case. The / in foo's definition makes an argument positional only so we can test them all by changing * to *args and adding **kwargs after d=5. \$\endgroup\$
    – Peilonrayz
    Dec 17, 2021 at 3:10

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.