How to typehint a function which takes a callable and its args as parameters - e.g. a numerical integration of a function?

user_6609896 0
Viewed 129

Problem

I have the following function (based on scipy.integrate.quad):

def simple_quad(func: Any, a: float, b: float, args: tuple = ()) -> float:
    def strips(n: int):
        for i in range(n):
            x = a + (b - a) * i / n
            yield func(x, *args) * 1 / n

    return sum(strips(1000))

... which basically evaluates func at a range of values and uses strips of fixed width to calculate the area under the graph. Optionally, parameters can be passed to func through the args tuple.

As you can see I've done some initial type-hints (in reality this is a .pyi stub for scipy), however I'm not happy with the typing of func & args being so loose. I want mypy to protect me from 2 things:

  1. func is a callable which must have a first positional parameter which is a float, and return a float, and can have optional *args
    • i.e. f(x:float, ...) -> float at the very least
    • I guess it can also have **kwargs (although can't have required name-only params or required positional params other than x)
  2. The optional positional *args to func must match the contents of the splatted args tuple

Examples

def cubic(x: float, a: float, b: float, c: float, d: float) -> float:
    "Function to evaluate a cubic polynomial with coefficients"
    return a + b * x + c * x ** 2 + d * x ** 3

simple_quad(cubic, a=1, b=10, args=(1, 2, 3, 4))  # fine, passes args to a,b,c,d and int is compatible with float signature
simple_quad(cubic, a=1, b=10)  # error from mypy as *args to `cubic` don't have default args so are all required
simple_quad(cubic, a=1, b=10, args=("a", "b", "c", "d"))  # arg length correct but type mismatch since cubic expects floats

x_squared: Callable[[float], float] = lambda x: x * x
simple_quad(x_squared, a=1, b=10, args=())  # should be fine as x_squared doesn't take any positional args other than x

def problematic(x: float, *, y: float) -> float: ...  # can't pass kwargs y via simple_quad, so can't be integrated

What I've tried

For func I tried to something with Protocol and Generics:

class OneDimensionalFunction(Protocol, Generic[T]): #double inheritance, although maybe I can get by with a metaclass for Generic
    def __call__(self, x: float, *args: T) -> float: ...

... in the hope I could write

def simple_quad(func: OneDimensionalFunction[T], a: float, b: float, args: tuple[T] = ()) -> float:

simple_quad(cubic, 1, 10, 10, (1,2,3,4)) # infer type requirement of args based on signature of func
# or
simple_quad[float,float,float,float](cubic, ...) #pass in additional type info above and beyond the expected Callable[[x:float,...], float]

... which has many problems I know, also Protocol doesn't play well with Callable if for example I want to pass in a lambda as func.

I labelled this python 3.10 as I think the new Parameter Specification Variables may help, but I've only seen those used in decorators so I'm not sure how to apply them here. Let me know your thoughts

python-3.10 mypy python generics
1 Answers

Use overloads with protocols. It's not great, but it gets you as close to real verification as you can probably get.

from typing import Protocol, Tuple, Any, overload, Union, Optional


class C1(Protocol):
    def __call__(self, a: float,                   ) -> float: ...
class C2(Protocol):
    def __call__(self, a: float, b: float,         ) -> float: ...
class C3(Protocol):
    def __call__(self, a: float, b: float, c: float) -> float: ...

ET = Tuple[()]
T1 = Tuple[float]
T2 = Tuple[float, float]
T3 = Tuple[float, float, float]

@overload
def quad(func: C1, a: float, b: float, args: Union[ET, T1] = ()) -> float: ...
@overload
def quad(func: C2, a: float, b: float, args: Union[ET, T2] = ()) -> float: ...
@overload
def quad(func: C3, a: float, b: float, args: Union[ET, T3] = ()) -> float: ...

def quad(func: Any, a: float, b: float, args: tuple = ()) -> float:
    return 0


def cubic(a: float, b: float, c: float, *, s: Optional[str] = None) -> float:
    return 0


quad(cubic, 0, 1)
user_3137680
Related