ParamSpec on overloaded & generic functions

Typing
1

Some thoughts on how ParamSpec interacts with generic and overloaded functions.

Generic functions

First example, with a generic function (works in mypy, pyright, ty, pyrefly, but not zuban):

from collections.abc import Callable

class CachedFunc[**P, R]:
    def __init__(self, inner: Callable[P, R]) -> None:
        self._inner: Callable[P, R] = inner
        self._cached: R | None = None

    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
        if self._cached is None:
            value = self._inner(*args, **kwargs)
            self._cached = value
        return self._cached

def identity[T](x: T, /) -> T:
    return x

fn = CachedFunc(identity)
#^ inferred as `CachedFunc[[T@identity], T@identity]`,
#  impossible to express as an actual annotation

x = fn(42)  # inferred as `int`
y = fn("banana")  # inferred as `str`, but is `int` at runtime

A CachedFunc remembers the first result if it exists. This is obviously problematic with an oveloaded or generic function, because the first return value might be wrong for the next call.

Overloads

Here is a more interesting example:

from typing import Literal, overload

def apple_sauce[**A, **B, AR, BR](
    _f: Callable[A, AR],
    _g: Callable[B, BR],
) -> Callable[A, BR]:
    raise NotImplementedError

@overload
def foo(x: int, y: str) -> Literal["1"]: ...
@overload
def foo(x: str, y: int) -> Literal["2"]: ...

@overload
def bar(a: str) -> Literal["3"]: ...
@overload
def bar(b: int) -> Literal["4"]: ...

criss_cross = apple_sauce(foo, bar)

the apple_sauce function uses the parameters of one function and the return type of another. This cruel science experiment leads to these results for different type checkers:

  • mypy and zuban infer the type of criss_cross as (x: int, y: str) -> Literal["3"]
  • ty uses the parameters of the first function, and a union of the second function’s returns as the return type:
    (x: int, y: str) -> Literal["3", "4"]
(x: str, y: int) -> Literal["3", "4"]
  • pyright produces a suspicious looking overload:
    Overload[
    (x: int, y: str) -> Literal['3'],
    (x: str, y: int) -> Literal['3'],
    (x: int, y: str) -> Literal['4'],
    (x: str, y: int) -> Literal['4']]
    in practice only the first two overload clauses will be used
  • pyrefly, when trying to call criss_cross, will show:
    Unexpected ParamSpec type: `[x: int, y: str] | [x: str, y: int]` [invalid-param-spec]
    which sounds more like a bug, since invalid-param-spec indicates that “ParamSpec is defined incorrectly or misused”, but here it’s inferred (or maybe the signature is supposed to be flagged instead?).

ty’s approach seems the most reasonable here. However, it does not work well in the “expected” case where the arguments are attached to the same function:

def identity[**P, T](fn: Callable[P, T]) -> Callable[P, T]:
    return fn

@overload
def f(x: int) -> int: ...
@overload
def f(x: str) -> str: ...
def f(x: object) -> object: return x

g = identity(f)
reveal_type(g)
#^ ty says: `Overload[(x: int) -> int | str, (x: str) -> int | str]`
#  but we'd really like to see the original overload, `Overload[(x: int) -> int, (x: str) -> str]`

New syntax&semantics?

I think these are symptoms of an issue with ParamSpec: it allows you to refer to the return type of a signature without relating it to the argument types. It ignores the ability of the return type to depend on argument types.

I don’t know what the syntax for it should be, but this is how I imagine a hypothetical alternative way to do what ParamSpec does, without the above issues:

# 1. For use cases where the return type may depend on the arguments:
def pass_through[F, **P: Params[F]](
    fn: F, /, *args: P.args, **kwargs: P.kwargs
) -> ReturnType[F, P]:  # or something like: P.result(F)?
    return fn(*args, **kwargs)

# 2. For use cases where the return type may not depend on the arguments:
def cached[**P, T](fn: Callable[P, T]) -> Callable[P, T]:
    ...

With this split:

  • the CachedFunc example will have the same type annotations. With generic and overloaded functions, generic and overloaded functions will be simplified: identity will be monomorphized as (object) -> object, and overloaded functions will have the return types of overloads unionized (as ty does) — using “option 2”
  • the apple_sauce example will also retain current type annotations, and it should have the same behaviour as in ty today — using “option 2”
  • Pass-through functions like asyncio.to_thread will use the way that lets them correctly wire up the return type with argument types — using “option 1”

One thing that this design affords is that you can now write the type for a pass-through decorator of a generic function:

class Logging[F: Callable]:
    def __init__(self, func: F) -> None:
        self._f = f

    def __call__[**P: Params[F]](
        self, /, *args: P.args, **kwargs: P.kwargs) -> ReturnType[F, P]: ...

logging_identity: Logging[IdentityFn] = Logging(identity)
# IdentityFn not included for brevity, but it's a Protocol with __call__

There is a difficulty with this design. Consider something like anyio.run’s signature: it needs to “destructure” the return type of the incoming callable. (the real signature uses *args only but let’s pretend it uses ParamSpec)

def run[**P, T](fn: Callable[P, Coroutine[T]]) -> T

I could think of something like this, inspired by Rust’s where clause

@Where[ReturnType[F, P]: Coroutine[T]]
def run[F, **P: Params[F], T](fn: F) -> T

That looks a bit wild. Maybe this use case could be skipped, and anyio.run could be limited to the “type 2” solution.
PEP 827 seems quite ambitious in its current state, but if something like this is implemented, then the return type could be marked as e.g. UnCoroutine[ReturnType[F, P]] where UnCoroutine is a user-defined thing.

Alternatives

If this feature is too complicated for the type system, or if it’s undesirable to add a second way to do a very similar thing, I think the current behaviour could be improved.

  • CachedFunc should monomorphize identity as (object) -> object. But pass-through cases should not. Maybe this can be changed depending on whether the R is stored as an attribute? (seems kinda implicit (also, the same issue is present if you implement cached_func as a closure with a nonlocal cached_result, that will be more difficult and in fact impossible to make a type stub for))
  • For apple_sauce, I think ty has a reasonable behaviour. (or it could be disallowed to mix signatures like that, since that is a purely synthetic example with no value)
8 Likes