Python Typing: Copy `**kwargs` from one function to another

Question

It is common pattern in Python extend or wrap functions and use **kwargs to pass all keyword arguments to the extended function.

i.e. take

class A:
    def bar(self, *, a: int, b: str, c: float) -> str:
       return f"{a}_{b}_{c}"
   

class B(A):
    def bar(self, **kwargs):
        return f"NEW_{super().bar(**kwargs)}"


def base_function(*, a: int, b: str, c: float) -> str:
    return f"{a}_{b}_{c}"


def extension(**kwargs) -> str:
    return f"NEW_{base_function(**kwargs)}"

Now calling extension(not_existing="a") or B().bar(not_existing="a") would lead to a TypeError, that could be detected by static type checkers.

How can I annotate my extension or B.bar in order to detect this problem before I run my code?

This annotation would be also helpful for IDE's to give me the correct suggestions for extension or B.bar.

Kound · Accepted Answer

Solution

Update: There is currently a CPython PR open to include the following solution into the standard library.

PEP 612 introduced the ParamSpec (see Documentation) Type.

We can exploit this to generate a decorator that tells our type checker, that the decorated functions has the same arguments as the given function:

from typing import (
   Callable, ParamSpec, TypeVar, cast, Any, Type, Literal,
)

# Define some specification, see documentation
P = ParamSpec("P")
T = TypeVar("T")

# For a help about decorator with parameters see 
# https://stackoverflow.com/questions/5929107/decorators-with-parameters
def copy_kwargs(
    kwargs_call: Callable[P, Any]
) -> Callable[[Callable[..., T]], Callable[P, T]]:
    """Decorator does nothing but returning the casted original function"""

    def return_func(func: Callable[..., T]) -> Callable[P, T]:
        return cast(Callable[P, T], func)

    return return_func

This will define a decorator than can be used to copy the complete ParameterSpec definition to our new function, keeping it's return value.

Let's test it (see also MyPy Playground)

# Our test function for kwargs
def source_func(foo: str, bar: int, default: bool = True) -> str:
    if not default:
        return "Not Default!"
    return f"{foo}_{bar}"

@copy_kwargs(source_func)
def kwargs_test(**kwargs) -> float:
    print(source_func(**kwargs))
    return 1.2

# define some expected return values
okay: float
broken_kwargs: float
broken_return: str

okay = kwargs_test(foo="a", bar=1)
broken_kwargs = kwargs_test(foo=1, bar="2")
broken_return = kwargs_test(foo="a", bar=1)

This works as expected with pyre 1.1.310, mypy 1.2.0 and PyCharm 2023.1.1. All three will complain about about the broken kwargs and broken return value. Only PyCharm has troubles to detect the default argument, as PEP 612 support is not yet fully implemented.

⚠️ Limitations

Still we need to by very careful how to apply this function. Assume the following call

runtime_error = kwargs_test("a", 1)

Will lead the runtime error “kwargs_test1() takes 0 positional arguments but 2 were given” without any type checker complaining.

So if you copy **kwargs like this, ensure that you put all positional arguments into your function. The function in which the parameters are defined should use keyword only arguments.

So a best practise source_func would look like this:

def source_func(*, foo: str, bar: int, default: bool = True) -> str:
    if not default:
        return "Not Default!"
    return f"{foo}_{bar}"

But as this is probably often used on library functions, we not always have control about the source_func, so keep this problem in mind!

You also could add *args to your target function to prevent this problem:

# Our test function for args and kwargs
def source_func_a(
  a: Literal["a"], b: Literal["b"], c: Literal["c"], d: Literal["d"], default: bool =True
) -> str:
    if not default:
        return "Not Default!"
    return f"{a}_{b}_{c};{d}"



@copy_kwargs(source_func_a)
def args_test(a: Literal["a"], *args, c: Literal["c"], **kwargs) -> float:
    kwargs["c"] = c
    # Note the correct types of source_func are not checked for kwargs and args,
    # if args_test doesn't define them (at least for mypy)
    print(source_func(a, *args, **kwargs))
    return 1.2

# define some expected return values
okay_args: float
okay_kwargs: float
broken_kwargs: float
broken_args: float

okay_args = args_test("a", "b", "c", "d")
okay_kwargs = args_test(a="a", b="b", c="c", d="d")
borken_args = args_test("not", "not", "not", "not")
broken_kwargs = args_test(a="not", b="not", c="not", d="not")

History of the PEP 612 Introduction for MyPy and PyCharm

MyPy and PyCharm had issues using ParamSpec when creating this answer. The issues seems to be resolved but the links are kept as historical reference:

MyPy merged a first implementation for ParamSpec on 7th April 2022
According to the related typedshed Issue, PyCharm should support ParamSpec but did not correctly detect the copied **kwargs but complained that okay = kwargs_test(foo="a", bar=1) would have invalid arguments. (Fixed now)
Mypy: Allow using TypedDict for more precise typing of **kwds
Mypy: PEP 612 tracking issue
Pyright Prototype support for typed **kwargs

Using Concatenate

If you want to copy the kwargs but also want to allow additional parameters you need to adopt kwargs with Concanate:

from typing import Concatenate

def copy_kwargs_with_int(
    kwargs_call: Callable[P, Any]
) -> Callable[[Callable[..., T]], Callable[Concatenate[int, P], T]]:
    """Decorator does nothing but returning the casted original function"""

    def return_func(func: Callable[..., T]) -> Callable[Concatenate[int, P], T]:
        return cast(Callable[Concatenate[int, P], T], func)

    return return_func

@copy_kwargs_with_int(source_func)
def something(first: int, *args, **kwargs) -> str:
    print(f"Yeah {first}")
    return str(source_func(*args, **kwargs))

something("a", "string", 3) # error: Argument 1 to "something" has incompatible type "str"; expected "int"  [arg-type]
okay_call: str
okay_call = something(3, "string", 3) # okay

See MyPy Play for details.

Note: Currently you need to define the a decorator for each variable you want to add and due to the nature of Concanate they can also just be added as args in front.

Python Typing: Copy `**kwargs` from one function to another

Answers (2)

Solution

⚠️ Limitations

History of the PEP 612 Introduction for MyPy and PyCharm

Using Concatenate

Related Questions