I often work with callables that have a lot of internal state (e.g., neural network layers), where the callable classes all inherit from a single base class. We can model this with ParamSpec:
from abc import abstractmethod
from typing import Generic, ParamSpec, TypeVar
T = TypeVar("T")
P = ParamSpec("P")
class BaseCallable(Generic[P, T]):
@abstractmethod
def apply(self, *args: P.args, **kwargs: P.kwargs) -> T:
"""Override this method."""
def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
# internal plumbing
return self.apply(*args, **kwargs)
However, we can only model relatively simple signatures like this:
class Identity(BaseCallable[[list[float]], list[float]]):
@override
def apply(self, input: list[float]) -> list[float]:
return input
iden = Identity()
reveal_type(iden([1.0, 2.0])) # list[float]
iden(input=[1.0, 2.0]) # type error: expected positional argument
but what if I want keyword arguments in my subclass of BaseCallable? The format for a concrete ParamSpec is of the form [int, str, bool]. There is no option to specify argument names.
But we can use Unpack and TypedDict from PEP 692 for this!
from random import random
from typing import NotRequired, TypedDict, Unpack
class Kwargs(TypedDict):
drop: NotRequired[bool] # `NotRequired` because it's an optional argument
class Dropout(BaseCallable[[list[float], Unpack[Kwargs]], list[float]]):
def __init__(self, probability: float) -> None:
assert 0 <= probability <= 1
self.prob = probability
def apply(self, input: list[float], drop: bool = True) -> list[float]:
if not drop:
return input
return [0.0 if random() < self.prob else x for x in input]
d = Dropout(probability=0.1)
d([1.0, 2.0]) # OK
d([1.0, 2.0], drop=False) # OK
So, the basic idea is to write concrete ParamSpecs as [int, bool, Unpack[SomeTypedDict]].
The meaning of Required, NotRequired and total is the same as was specified in PEP 692, and allows us to specify required and optional keyword arguments.
There are still some kinds of signatures that cannot be expressed this way, but this is much better than before! And it doesn’t require any new syntax or any new code in the standard library.
(This post was prompted by this recent mypy PR: Allow TypedDict unpacking in Callable types by ilevkivskyi · Pull Request #16083 · python/mypy · GitHub )