Migwell
Migwell

Reputation: 20107

How to annotate the type of arguments forwarded to another function?

Let's say we have a trivial function that calls open() but with a fixed argument:

def open_for_writing(*args, **kwargs):
    kwargs['mode'] = 'w'
    return open(*args, **kwargs)

If I now try to call open_for_writing(some_fake_arg = 123), no type checker (e.g. mypy) can tell that this is an incorrect invocation: it's missing the required file argument, and is adding another argument that isn't part of the open signature.

How can I tell the type checker that *args and **kwargs must be a subset of the open parameter spec? I realise Python 3.10 has the new ParamSpec type, but it doesn't seem to apply here because you can't get the ParamSpec of a concrete function like open.

Upvotes: 19

Views: 2600

Answers (2)

InSync
InSync

Reputation: 10492

@SimonHawe's answer is correct. However, it could be simplified and modified a little bit to avoid --strict rants.

  1. Use Any for the target function's *args, **kwargs and return type. They will be overwritten by our decorator anyway.
  2. With cast(), the target function can just be returned as-is, hence no wraps() is needed.

(playgrounds: Mypy, Pyright)

from collections.abc import Callable
from typing import Any, cast, ParamSpec, TypeVar

T = TypeVar('T')
P = ParamSpec('P')

def copy_signature_from(_origin: Callable[P, T]) -> Callable[[Callable[..., Any]], Callable[P, T]]:
    def decorator(target: Callable[..., Any]) -> Callable[P, T]:
        return cast(Callable[P, T], target)
    
    return decorator

Example usage:

def f(foo: str, bar: int, *, baz: bool) -> None:
    print(foo, bar, baz)

@copy_signature_from(f)
def f_wrapper(*args: Any, **kwargs: Any) -> Any:
    return f(*args, **kwargs)
class Parent:
    def __init__(self, foo: str, bar: int, *, baz: bool) -> None:
        ...

class Child(Parent):
    @copy_signature_from(Parent.__init__)
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        ...

There used to be one small caveat with Pyright; it would emit an error if the origin function is @overload-ed. This is a bug and it was fixed in 1.1.351.

from typing import overload, Literal

@overload
def g(v: Literal[1]) -> int:
    ...
    
@overload
def g(v: Literal[2]) -> str:
    ...

def g(v: Literal[1, 2]) -> int | str:
    ...

@copy_signature_from(g)
def g_wrapper(*args: Any, **kwargs: Any) -> Any:
    # mypy    => fine
    # pyright => error: "g_wrapper" is marked as overload, but no implementation is provided
    ...

IDE support

VSCode/Pylance will provide good parameter autocompletions:

(foo: str, bar: int, *, baz: bool) -> None

...and so does PyCharm:

f(foo: str, bar: int, *, baz: bool)

Upvotes: 4

Simon Hawe
Simon Hawe

Reputation: 4529

I think out of the box this is not possible. However, you could write a decorator that takes the function that contains the arguments you want to get checked for (open in your case) as an input and returns the decorated function, i.e. open_for_writing in your case. This of course only works with python 3.10 or using typing_extensions as it makes use of ParamSpec

from typing import TypeVar, ParamSpec, Callable, Optional

T = TypeVar('T')
P = ParamSpec('P')


def take_annotation_from(this: Callable[P, Optional[T]]) -> Callable[[Callable], Callable[P, Optional[T]]]:
    def decorator(real_function: Callable) -> Callable[P, Optional[T]]:
        def new_function(*args: P.args, **kwargs: P.kwargs) -> Optional[T]:
            return real_function(*args, **kwargs)

        return new_function
    return decorator

@take_annotation_from(open)
def open_for_writing(*args, **kwargs):
    kwargs['mode'] = 'w'
    return open(*args, **kwargs)


open_for_writing(some_fake_arg=123)
open_for_writing(file='')

As shown here, mypy complains now about getting an unknown argument.

Upvotes: 11

Related Questions