overload and contextmanager doesn't work together

[prousso]

Bug Report

contextmanager and overload doesn't seem to be compatbile together

To Reproduce

from contextlib import contextmanager
from typing import overload, ContextManager, Iterator
from uuid import UUID

@overload
def cactus() -> ContextManager[int]:
    ...

@overload
def cactus(_id: str) -> ContextManager[UUID]:
    ...

@contextmanager
def cactus(_id: str | None = None) -> Iterator[UUID] | Iterator[int]:
    if not _id:
        yield 1
    yield UUID(_id)

Expected Behavior

Mypy should find no issue

Actual Behavior

test.py: note: In function "cactus":
test.py:16:2: error: Overloaded function implementation cannot produce return type of signature 1  [misc]
    @contextmanager
     ^
test.py:16:2: error: Overloaded function implementation cannot produce return type of signature 2  [misc]
    @contextmanager
     ^
Found 2 errors in 1 file (checked 1 source file)

Your Environment

• Mypy version used: 1.0.0
• Mypy command-line flags:
• Mypy configuration options from mypy.ini (and other config files):
• Python version used: 3.10.6

[donsokolone]

I second that, ran into this with Python 3.9.16 and mypy 1.3.0.

[ikonst]

Appears to reproduce without context managers:
https://mypy-play.net/?mypy=latest&python=3.11&gist=970a3308fdf33f864002c17c4847f53e

[tyralla]

I also ran into the same problem today. My initial thoughts:

First, the reported errors in the example of @ikonst seem to result from ignoring the parameter name "x" in the wrapper's definition. This can be fixed by using Protocol instead of Callable:

from typing import Callable, overload, Protocol

class P(Protocol):
    def __call__(self, a: int | str) -> int | str: ...

def wrap(f: P) -> P: ...

@overload
def f(a: int) -> int: ...
@overload
def f(a: str) -> str: ...
@wrap
def f(a: int | str) -> int | str: ...

reveal_type(f(1))  # note: Revealed type is "builtins.int"

Or, more generally, by using ParamSpec:

from typing import Callable, overload, ParamSpec

P = ParamSpec("P")

def wrap(f: Callable[P, int | str]) -> Callable[P, int | str]: ...

@overload
def f(a: int) -> int: ...
@overload
def f(a: str) -> str: ...
@wrap
def f(a: int | str) -> int | str: ...

reveal_type(f(1))  # note: Revealed type is "builtins.int"

Second, regarding the initial report by @prousso, I looked into the class hierarchy. contextmanager returns Callable[_P, _GeneratorContextManager[_T_co]]. _GeneratorContextManager is a subclass of AbstractContextManager and ContextDecorator. ContextManager is an alias for AbstractContextManager. So, in the initial report, the implementation return type (_GeneratorContextManager) is a subtype of the overloaded return type (ContextManager).

Always returning _GeneratorContextManager circumvents the error reports:

from contextlib import contextmanager, _GeneratorContextManager
from typing import overload, Iterator
from uuid import UUID

@overload
def cactus() -> _GeneratorContextManager[int]: ...
@overload
def cactus(_id: str) -> _GeneratorContextManager[UUID]: ...
@contextmanager
def cactus(_id: str | None = None) -> Iterator[UUID] | Iterator[int]: ...

So, the question is whether Mypy should allow the implementation to return subtypes of what is defined by the overloads. Here is a simple example of what's currently possibly going wrong:

from typing import overload

class A: ...
class B: ...
class AA(A): ...

@overload
def f(x: A) -> A: ...
@overload
def f(x: B) -> B: ...
def f(x: A | B) -> AA | B: ...  # E: Overloaded function implementation cannot produce return type of signature 1

I cannot see a reason for disallowing this. I could try to fix this in the following days, but I would first want approval that we are dealing with a Mypy limitation here, not with expected behaviour. @JukkaL: Could you or someone else from the core team have a quick look, please?

[ilevkivskyi]

Uniform application of @contextmanager fixes the issue (on current master):

from contextlib import contextmanager
from typing import overload, ContextManager, Iterator
from uuid import UUID

@overload
@contextmanager
def cactus() -> Iterator[int]:
    ...

@overload
@contextmanager
def cactus(_id: str) -> Iterator[UUID]:
    ...

@contextmanager
def cactus(_id: str | None = None) -> Iterator[UUID] | Iterator[int]:
    if not _id:
        yield 1
    yield UUID(_id)

Btw I think we should actually recommend this approach for decorated overloads (after/during the next release). Having the decorator "manually" applied to each overload by the user may be error-prone (OTOH one will see the already decorated signature, but IMO it is not a big advantage).

[tyralla]

It's very nice that Mypy will soon understand decorated overloads. Thanks for this fix!

Your change does not address the cause of why Mypy emits a false alarm for the original example code, does it? (But maybe it does not matter anymore because no one following the new recommendation will likely let the implementation return subtypes of the overloads' return types.)

[ilevkivskyi]

@tyralla
The original error is non-trivial to "fix". Because it is not even clear what is the desired state. Current sanity check for return types is very simple: overload item return type must be either a subtype or a supertype of the return type of implementation (i.e. they must be somehow related). In your example with AA neither is the case because A is neither a subtype of AA | B nor a supertype of it. Special-casing unions is technically possible, but would be quite ugly, and value will be quite limited.

[tyralla]

Ah, this explains why Mypy fails here. Thank you. I strongly assumed that Mypy would already check the individual members of the Union returned by the implementation (if it returns a Union). This seemed natural to me because the implementation often just collects the return types of the overloads within a Union.

I do not entirely understand what "desired state" you are writing of. But if you are afraid, there would be too much code complication for fixing a seldom noticeable limitation (thanks to your elegant workaround), closing this issue (as half-solved and half-wont-be-fixed) seems right.

[tyralla]

During a (nice summerly) training run, it came to my mind that, following @ilevkivskyi's explanation, Mypy would not detect the following bug:

from typing import overload

class A: ...
class B: ...
class AA(A): x: int = 0

@overload
def f(x: A) -> AA: ...
@overload
def f(x: B) -> B: ...
def f(x: A | B) -> A | B: return x

f(A()).x    # crash but no error report

So, the problem reported by @prousse seems to be strongly related to #12401, where the "Union compromise" was also suggested by @KotlinIsland.

However, there is also a discussion in issue #15790 regarding the possible future explicit consideration of intersections, which might be the unclear desired state @ilevkivskyi meant.