Search code examples
pythonpython-typing

Can ParamSpec be used to type individual arguments?


I'm trying to type the following wrapper function. It takes another function and the function's arguments and runs it with some side effect.

from typing import Callable, ParamSpec, TypeVar

P = ParamSpec("P")
R = TypeVar("R")


def wrapper(func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
    # Run some side effects
    return func(*args, **kwargs)


def f(a: int, b: int) -> int:
    return a + b

# Prints 3 with the side effect. 
print(wrapper(f, 1, 2))

The above snippet works and mypy is happy. However, I don't like the black boxes that *args and **kwargs are and would like to make wrapper less generic. How do I type something similar to this:

from typing import Callable, ParamSpec, TypeVar

P = ParamSpec("P")
R = TypeVar("R")

# How do I express the fact that a and b are the input params of func?
def wrapper(func: Callable[P, R], a: ?, b: ?) -> R:
    # Run some side effects
    return func(a, b)


def f(a: int, b: int) -> int:
    return a + b

# Mypy should be happy here!
print(wrapper(f, 1, 2)) 

Is this possible?


Solution

  • When you know the arguments that will be passed to func, you don't need ParamSpec; ordinary TypeVars will do.

    AType = TypeVar('AType')
    BType = TypeVar('BType')
    
    
    def wrapper(func: Callable[[AType, BType], R], a: AType, b: BType) -> R:
        return func(a, b)
    

    If there were additional arguments, then you would use Concatenate to combine the known arguments with ParamSpec.

    AType = TypeVar('AType')
    BType = TypeVar('BType')
    P = ParamSpec('P')
    
    
    def wrapper(func: Callable[Concatenate[AType, BType, P], R],
                a: AType,
                b: BType,
                *args: P.args,
                **kwargs: P.kwargs) -> R:
        return func(a, b, *args, **kwargs)