ParamSpec does not consider substitution of @overloads

[sbrudenell]

Bug Report

It looks like mypy only considers the first @overload of a function, when considering substitutions for a ParamSpec.

In particular this affects use of asyncio.to_thread(f, arg), if f has @overloads.

To Reproduce

import asyncio
from collections.abc import Callable
from typing import overload
from typing import TypeVar

from typing_extensions import ParamSpec


@overload
def f(x: int) -> None:
    ...


@overload
def f(x: str) -> None:
    ...


def f(x) -> None:
    pass


_P = ParamSpec("_P")
_R = TypeVar("_R")


# type signature modeled after asyncio.to_thread()
def generic(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs) -> _R:
    ...


# using first overload: passes
generic(f, 1)
# using second overload: fails
generic(f, "test")

Expected Behavior

$ mypy t.py
Success: no issues found in 1 source file

Actual Behavior

$ mypy t.py
t.py:35: error: Argument 2 to "generic" has incompatible type "str"; expected "int"
Found 1 error in 1 file (checked 1 source file)

Your Environment

• Mypy version used: mypy 0.971 (compiled: yes)
• Mypy command-line flags: mypy t.py
• Mypy configuration options from mypy.ini (and other config files): none
• Python version used: 3.10.4
• Operating system and version: Ubuntu 22.04.1

[erictraut]

PEP 612 unfortunately didn't contemplate how ParamSpec would work with overloads, so different type checkers treat it differently. Pyre attempts to capture all overloaded signatures in the ParamSpec and then produce an overloaded result when calculating the specialized return type. This approach works for common cases but falls apart at the edges. Mypy chooses the first overload. Once again, this works for some cases but is probably not what was intended in many cases. Pyright currently disallows the use of overloads with ParamSpec (emits an error) since the behavior is not well defined.

If this is a problem that's worth addressing (and I'm not convinced it is), then someone will probably need to work out all of the intended behaviors (including edge cases), document everything, and drive consensus among the various type checker maintainers.

[sbrudenell]

Could you elaborate on what the edge cases are, for the interaction of ParamSpec and @overload? They're not obvious to me and I couldn't find any discussion of this on the various forums.

(I did find you mentioned this interaction in #12292 (comment), but not that there were unresolved issues and edge cases)

It seems worthwhile to me to address this, since I hit the bug 😄. But perhaps I don't understand the tradeoffs since I don't understand the edge cases.

[erictraut]

To identify the edge cases, you will probably need to write a prototype implementation and work through a bunch of test cases that include ParamSpec and Concatenate and a bunch of overloads (including those with function-scoped type variables and ParamSpecs).

I think the approach will necessarily involve solving for not just one set of type variables (as is done with a normal call expression) but n sets, one for each overload in the argument passed to the call. And then once you've solved all n sets of type variables, you'll need to figure out how to combine the results to get the final return type of the call. If the declared return type is a callable itself that uses the ParamSpec, it will presumably translate into a new synthesized overload. If the declared return type isn't a callable, then you will presumably need to apply a union or join from all of the solved sets of type variables to get a combined answer.

It's even worse if the target of the call accepts two callable input parameters that are parameterized by separate ParamSpecs. In that case, if you want to support passing overloaded functions for both parameters, you will need to solve n x m sets of type variables where n is the number of overloads in the first argument and m is the number of overloads in the second.

That's just one issue I can think of, but I'm sure there are many more that I haven't.

[valsteen]

I've met this issue while implementing something very similar as the issue description, using starlette's run_in_threadpool which uses ParamSpec. This didn't play well with boto3-stubs which uses @overload to describe the different flavours of boto3.resource. But there is a way around using a lambda closure. This passes:

generic(lambda: f("test"))

and this plays nice because mypy correctly infers the return type of a lambda. Here is a modification of the example with lambda, returning a value and using assert_type to demonstrate that the type is preserved:

https://mypy-play.net/?mypy=latest&python=3.10&gist=1f323faa2788a034a665dd4052d1e51a

from collections.abc import Callable
from typing import Any, ParamSpec, TypeVar, overload

from typing_extensions import assert_type


@overload
def f(x: int) -> list[int]:
    ...


@overload
def f(x: str) -> list[str]:
    ...


def f(x):
    return [x]


_P = ParamSpec("_P")
_R = TypeVar("_R")


def generic(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs) -> _R:
    return func(*args, **kwargs)


list_int = generic(lambda: f(1))
assert_type(list_int, list[int])
list_str = generic(lambda: f("test"))
assert_type(list_str, list[str])

[erictraut]

FWIW, I worked through all of the issues with ParamSpec capturing an overloaded signature and implemented the support in pyright. There are some limitations, but it works well for all of the real-world use cases that I've encountered. There are several things to keep in mind:

1. The constraint solver needs to solve constraints independently for each of the applicable overloads. This required a bunch of refactoring in pyright. I'm not sure how much work that would be in mypy.
2. Overloads need to be filtered based on the signature of the accepting function. For example, if the parameter type is Callable[Concatente[int, P], None] and the argument is an overloaded callable, any overload whose first parameter is not int is not applicable and should be eliminated. Likewise, if the accepting function accepts *args: P.args and **kwargs: P.kwargs, then the overload list should be filtered based on the supplied arguments that map to *args and **kwargs. If this filtering process eliminates all of the overloads, then a type error should be generated. If a single overload remains after filtering, then it's treated like a non-overloaded function. If more than one overload remains, then each overloaded signature gets its own constraint solver context.
3. If the receiving call returns a callable type that includes the ParamSpec (which is often the case), the result is a new (synthesized) overload composed of all of the individual constraint solutions. If the call returns a non-callable type, then it's not clear which constraint solution should be used. In this case, pyright picks the first constraint solution (i.e. for the first overload in the filtered list produced from step 2). It might also be reasonable to return a union or join of all the solutions. In real-world code, I don't think this ever happens.

Note: Things get very complicated when more than one ParamSpec is involved in the signature because you need to create a constraint solver context for all combinations of ParamSpecs, but this is extremely rare, so it can probably be ignored or treated as an error.

If someone is interested in adding this support to mypy, feel free to reuse the unit tests I wrote for pyright.