PEP 677 with an easier-to-parse and more expressive syntax

Typing
27

The discussion here about making FunctionType generic, like Callable is today, made me consider this idea again.

Here is how we could have a better Callable, and also a better FunctionType, and any other kind of callable that has properties.

There would be a new TypeVar-like called SignatureSpec (name can still be bikeshedded):

F = SignatureSpec("F")

It works similar to ParamSpec, except that it also includes the return type and has a specialized syntax for specifying it (as we will see). Its PEP 695 syntax is F() (can still be bikeshedded) and it’s used like this:

class FunctionType[F()]:
    __name__: str
    __qualname__: str
	def __call__(self, *args: F.args, **kwargs: F.kwargs) -> F.returns: ...

For specifying the SignatureSpec, the Extended Syntax Supporting Named and Optional Arguments from PEP 677 is used:

type MyFunc = FunctionType[(int, y: float, *, z: bool = ...) -> bool]

Note that, in contrast to PEP 677, the (...) -> ... syntax can now only appear within square brackets. This hopefully makes parsing (by humans and machines) less expensive.

Unfortunately, we cannot change the definition of Callable to use SignatureSpec without breaking tons of already existing typed code. But of course we would want to use the new syntax for Callables, so what to do?

One option is to introduce a sort of type alias for Callable, which uses SignatureSpec. It could be called typing.Fn (name not final), such that

from collections.abc import Callable

def reduce(
    func: Callable[[int, int], int],
    l: list[int],
) -> int: ...

becomes

from typing import Fn

def reduce(
    func: Fn[(int, int) -> int],
    l: list[int],
) -> int: ...

One of the nice things about this approach is that it allows the specification of “complex callables” — which are not uncommon in Python — with a natural-looking syntax. We have already seen FunctionType, but there is more.

Take for example nn.Module from PyTorch. What PyTorch calls modules are basically callables but they also carry state that is updated during gradient descent. The user of the library sub-classes nn.Module and implements the forward() method which is then called by __call__() (which also does other things needed for computing gradients and things like that).

With SignatureSpec, nn.Module could be defined like this:

class Module[F()]:
	@abstractmethod
    def forward(self, *args: F.args, **kwargs: F.kwargs) -> F.returns: ...
    def __call__(self, *args: F.args, **kwargs: F.kwargs) -> F.returns:
        # do other stuff here
        return self.forward(*args, **kwargs)

A precise signature of a PyTorch module can then be specified with:

layer: nn.Module[(Tensor, my_flag: bool) -> Tensor]
layer(torch.zeros(4), my_flag=False)

What is described here is to some degree already possible with ParamSpec, but note for example that one cannot specify the name of the boolean flag with ParamSpec.

nn.Module can be sub-classed like this (example from here):

type Activations = Literal["relu", "lrelu"]
class DynamicNet(nn.Module[(Tensor, act: Activations) -> Tensor]):
    def __init__(self, num_layers: int):
        super().__init__()
        self.linears = nn.ModuleList(
            [MyLinear(4, 4) for _ in range(num_layers)])
        self.activations = nn.ModuleDict({
            'relu': nn.ReLU(),
            'lrelu': nn.LeakyReLU(),
        })
    def forward(self, x: Tensor, act: Activations) -> Tensor:
        for linear in self.linears:
            x = self.activations[act](linear(x))
        return x

net = DynamicNet(2)
net(torch.ones(9), act="sigmoid")  # type checker: "sigmoid" is not allowed here

We have to specify the signature here twice (once in nn.Module[] and once in the definition of forward()), but the type checker will inform us when they get out of sync, so this doesn’t seem so bad.

One criticism of PEP 677 was that nested callables start to look very confusing. Consider this example from the PEP:

def f() -> (int) -> (str) -> bool: pass

With this proposal here, it would be

def f() -> Fn[(int) -> Fn[(str) -> bool]]: pass

which is less confusing, I would say.

There are some clear disadvantages to this approach:

  • It’s still required to import something (typing.Fn) to specify callables. One selling point of PEP 677 was that the import of Callable wasn’t necessary anymore, but the import will unfortunately still be necessary with this proposal.
  • I think it isn’t possible to specify generic callables (that is, callables that have their own bound TypeVar as in lambda[T](T) -> T), but that’s also not possible with today’s Callable and also wasn’t possible with PEP 677.
  • The lexer still needs to perform a look-ahead in order to recognize the syntax, but at least it’s only needed within square brackets.

But the advantage is that this approach can more accurately describe the breadth of callables that are used in Python.

1 Like