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:
    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):
    return f"NEW_{super().bar(**kwargs)}"

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

How can I annotate my extension 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.

Answer 1

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.

⚠️ 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")

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

Answer 2

Based on @kound answer.

To remain DRY, we can do the same without re-declaring return type. Type variable T will be deduced later (not when copy_kwargs is called, but when its returned function is), but it doesn't affect further type checking.

from typing import Callable, ParamSpec, TypeVar, cast, Any

# Our test function
def source_func(*, foo: str, bar: int) -> str:
    return f"{foo}_{bar}"

# 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


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

reveal_type(kwargs_test(foo="a", bar=1))
reveal_type(kwargs_test(foo=1, bar="2"))

And here's mypy playground link to look at this in action.