Wrapped overloaded function drops all but the first overload

[maffoo]

Bug Report

When wrapping an overloaded function and using ParamSpec to preserve the call signature, only the first overload is preserved on the resulting function.

To Reproduce

See:

• https://gist.github.com/mypy-play/96895b9df7aa87c9e9f6cc59203d0751
• https://mypy-play.net/?mypy=latest&python=3.11&gist=96895b9df7aa87c9e9f6cc59203d0751

from typing import Callable, overload, TypeVar
from typing_extensions import ParamSpec

P = ParamSpec("P")
T = TypeVar("T")

def transform(func: Callable[P, list[T]]) -> Callable[P, T]:
    def wrapped(*args: P.args, **kwargs: P.kwargs) -> T:
        return func(*args, **kwargs)[0]

    return wrapped

@overload
def foo(x: int) -> list[float]:
    pass

@overload
def foo(x: str) -> list[str]:
    pass

def foo(x: int | str) -> list[float] | list[str]:
    if isinstance(x, int):
        return [1 / x]
    return [x[::-1]]

bar = transform(foo)

assert bar(2) == 0.5
assert bar("abc") == "cba"

reveal_type(foo)
reveal_type(bar)

Expected Behavior

I would expect the overloads to be preserved on the decorated function, e.g.

main.py:31: note: Revealed type is "Overload(def (x: builtins.int) -> builtins.list[builtins.float], def (x: builtins.str) -> builtins.list[builtins.str])"
main.py:32: note: Revealed type is Overload(def (x: builtins.int) -> builtins.float, def (x: builtins.str) -> builtins.str)"

Actual Behavior

main.py:29: error: Argument 1 has incompatible type "str"; expected "int"  [arg-type]
main.py:31: note: Revealed type is "Overload(def (x: builtins.int) -> builtins.list[builtins.float], def (x: builtins.str) -> builtins.list[builtins.str])"
main.py:32: note: Revealed type is "def (x: builtins.int) -> builtins.float"

Your Environment

• Mypy version used: 1.4.1
• Python version used: 3.10.10

[erictraut]

Supporting overloads for ParamSpec is difficult. PEP 612 (which introduced ParamSpec) didn't contemplate this case or discuss the behaviors and limitations thereof. If the ParamSpec is used to parameterize a class or a generic type alias, it's not clear how to handle an overload. However, in the case where the function returns a Callable[P, T] (as the example above does), the meaning and behavior is well defined. I've implemented support for this in pyright, so the above type checks as you would expect. If someone is interested in implementing this in mypy and would like some tips or suggestions based on my experience, let me know.

[hauntsaninja]

@maffoo is the project you were running into this on open source?

[maffoo]

@maffoo is the project you were running into this on open source?

Hi @hauntsaninja, yes, the project where I ran into this is in "duet", a small library for doing concurrency with futures: https://github.com/google/duet. In particular, there's a function duet.sync that wraps an async function into a synchronous version and we've been using a mypy plugin to get the typing right (https://github.com/google/duet/blob/main/duet/typing.py), but really what we'd like is to use ParamSpec and define duet.sync like:

def sync(f: Callable[P, Awaitable[T]]) -> Callable[P, T]:
    ...

Unfortunately, we use duet.sync internally with overloaded functions in many places, which is how I encountered this issue. I have a branch where I've switched to using ParamSpec (https://github.com/google/duet/tree/u/maffoo/param-spec). I tried out the latest version of mypy with that branch and it seems to work as expected (in particular see the file https://github.com/google/duet/blob/u/maffoo/param-spec/check.py). Thanks for fixing this!