Unpacking TypeDicts for specifying more complex ParamSpecs

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]):
    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]]):
    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 )

Great thinking! I proposed this here if you want to review the small discussion.

Perhaps there’s enough here to write up a PEP?