mypy: How to apply *args on a Callable type?

Viewed 69

I have a function like below,

def run_the_f(f):
  # run f function after some validation

Based on some condition, the signature of f function changes like f(1.0), f(1.0,2.0), f(1.0,2.0,3.0,..). In other words, the number of input arguments can be varying in f similar to udf f in pyspark.

I am using mypy and I tried below which is failing,

def run_the_f(f: Callable[[*float],float]):
  # run f after some validation

Can someone support in what to fill in the Callable?

1 Answers

Using typing.Protocol:

from typing import Protocol, TypeVar
from abc import abstractmethod

T_contra = TypeVar('T_contra', contravariant=True)
T_co = TypeVar('T_co', covariant=True)


class SupportsCallWithArgs(Protocol[T_contra, T_co]):

    @abstractmethod
    def __call__(self, *args: T_contra) -> T_co:
        pass


def run_the_f(f: SupportsCallWithArgs[float, float]):
    f(1., 2., 3.)   # pass
    f([])           # fail: Argument to "__call__" of "SupportsCallWithArgs" has incompatible type "List[<nothing>]";
                    # expected "float"

def f(*args: float) -> float:
    pass


run_the_f(f)             # pass
run_the_f(lambda: None)  # fail: Argument 1 to "run_the_f" has incompatible type "Callable[[], None]";
                         # expected "SupportsCallWithArgs[float, float]"

VarArg in mypy's expansion package maybe a simpler choice (according to the reference provided in the comment area, this choice is not recommended):

from collections.abc import Callable
from mypy_extensions import VarArg


def run_the_f(f: Callable[[VarArg(float)], float]):
    f(1., 2., 3.)    # pass

But this seems to be just mypy's self hypnosis, Pycharm's type checker will prompt 2. and 3. as unexpected arguments of f:

# mypy_extensions.py
def VarArg(type=Any):
    """A *args-style variadic positional argument"""
    return type
Related