Generic Keyword Arguments Type Annotations

157 Views Asked by At

I would like to add type annotations for fn1, fn2 and kwargs:

class Foo:
    def __init__(self, fn1, fn2, **kwargs):
        self.res1 = fn1(**kwargs)
        self.res2 = fn2(**kwargs)

I want to bound kwargs to fn1 & fn2 input arguments to avoid cases like:

# fails on fn2 as kwargs doesnt have a key y
Foo(fn1=lambda x: x+1, fn2=lambda x,y:x+y, x=1).

Is it even possible to annotate it correctly in current Python (3.11 or 3.12)?

For example, if i had only a single parameter to the functions, I'd use typing.TypeVar as follows:

from typing import Callable, TypeVar

T = TypeVar('T')


class Foo:
    def __init__(self, fn1: Callable[[T], T], fn2: Callable[[T], T], x: T):
        self.res1 = fn1(x)
        self.res2 = fn2(x)


# mypy typecheck fail when running mypy file.py
Foo(lambda x, y: x+1, lambda y: y*2, x=5)

# mypy typecheck pass when running mypy file.py
Foo(lambda x: x+x, lambda y: y*2, x='a')
Foo(lambda x: x+x, lambda y: y*2, x=3.14)

mypy file.py output:

file.py:13: error: Cannot infer type of lambda  [misc]
file.py:13: error: Argument 1 to "Foo" has incompatible type "Callable[[Any, Any], Any]"; expected "Callable[[int], int]"  [arg-type]
Found 2 errors in 1 file (checked 1 source file)
1

There are 1 best solutions below

0
Wombatz On

This can be nearly done with ParamSpec. There is a small problem however, to use ParamSpec you must also allow *args for your functions. You cannot separate kwargs and args with ParamSpec.

from typing import ParamSpec, TypeVar

from collections.abc import Callable

# the return type
R = TypeVar('R')
# param type
P = ParamSpec('P')

class Foo:
    # you need *args
    def __init__(
        self,
        fn1: Callable[P, R],
        fn2: Callable[P, R],
        *args: P.args,
        **kwargs: P.kwargs,
    ) -> None:
        self.res1 = fn1(*args, **kwargs)
        self.res2 = fn2(*args, **kwargs)

Now this works:

Foo(
    fn1=lambda x: x+1,
    fn2=lambda x: x+2,
    x=1,
)

This produces a nice error message:

Unexpected keyword argument "y" for "Foo"

Foo(
    fn1=lambda x: x+1,
    fn2=lambda x: x+2,
    y=1,
)

And this produces a pretty complex error message

Foo(
    fn1=lambda x: x+1,
    fn2=lambda y: y+2,
    y=1,
)