Wrong return type from ExitStack.enter_context when given a child of ExitStack

[thatguystone]

Bug Report

The return type from enter_context() is incorrect when the argument is a child of ExitStack.

To Reproduce

# ctx.py

from contextlib import ExitStack

class Thing(ExitStack):
    def foo(self) -> None:
        print("foo")

stack = ExitStack()
thing = stack.enter_context(Thing())
reveal_type(thing)
thing.foo()

Expected Behavior

The type of thing should be Thing, not ExitStack.

Actual Behavior

$ mypy ctx.py 
ctx.py:11: note: Revealed type is "contextlib.ExitStack"
ctx.py:12: error: "ExitStack" has no attribute "foo"
Found 1 error in 1 file (checked 1 source file)

Your Environment

$ mypy --version 
mypy 0.960 (compiled: yes)

[AlexWaygood]

The typeshed stubs look correct here to me — @JelleZijlstra, are you sure this is a typeshed bug and not a mypy bug?

[JelleZijlstra]

This is because of this line:

typeshed/stdlib/contextlib.pyi

Line 166 in 62a8a69

class ExitStack(AbstractContextManager[ExitStack]):

. ExitStack inherits from AbstractContextManager[ExitStack]. Logically this should be AbstractContextManager[Self], but that's not allowed. As a result, even though __enter__ is annotated as returning Self, when we match a subclass of ExitStack against AbstractContextManager[T], T gets solved to ExitStack instead of the subclass.

Here is a simple repro with mypy:

from typing import TypeVar
from contextlib import AbstractContextManager

T = TypeVar("T")
Self = TypeVar("Self")

class ExitStack(AbstractContextManager[ExitStack]):
    def __enter__(self: Self) -> Self: ...
    def __exit__(self, *args: object) -> bool: ...

class A(ExitStack):
    pass

def f(x: AbstractContextManager[T]) -> T:
    raise

reveal_type(f(ExitStack()))  # Revealed type is "__main__.ExitStack"
reveal_type(f(A()))  # Revealed type is "__main__.ExitStack"

One solution is to define our own protocol that ExitStack doesn't inherit from explicitly:

from typing import TypeVar, Protocol
from contextlib import AbstractContextManager

T = TypeVar("T")
T_co = TypeVar("T_co", covariant=True)
Self = TypeVar("Self")

class _CMProto(Protocol[T_co]):
    def __enter__(self) -> T_co: ...
    def __exit__(self, *args: object) -> bool | None: ...

class ExitStack(AbstractContextManager[ExitStack]):
    def __enter__(self: Self) -> Self: ...
    def __exit__(self, *args: object) -> bool: ...

class A(ExitStack):
    pass

def f(x: _CMProto[T]) -> T:
    raise

reveal_type(f(ExitStack()))  # Revealed type is "__main__.ExitStack"
reveal_type(f(A())) # Revealed type is "__main__.A"

[AlexWaygood]

Not having ExitStack inherit from anything also seems to solve the problem:

from typing import TypeVar
from contextlib import AbstractContextManager

T = TypeVar("T")
Self = TypeVar("Self")

class ExitStack:
    def __enter__(self: Self) -> Self: ...
    def __exit__(self, *args: object) -> bool: ...

class A(ExitStack):
    pass

def f(x: AbstractContextManager[T]) -> T:
    raise

reveal_type(f(ExitStack()))  # Revealed type is "__main__.ExitStack"
reveal_type(f(A()))  # Revealed type is "__main__.A"

Mypy still understands that ExitStack is a subtype of AbstractContextManager since it matches the structure of the protocol, but it can now infer the type properly.

https://mypy-play.net/?mypy=latest&python=3.10&gist=615a2b644f87ffc9f42397f83853ed17

[JelleZijlstra]

Not having ExitStack inherit from anything also seems to solve the problem:

But ExitStack inherits from AbstractContextManager at runtime, so ideally the stubs should reflect that.

[AlexWaygood]

Not having ExitStack inherit from anything also seems to solve the problem:

But ExitStack inherits from AbstractContextManager at runtime, so ideally the stubs should reflect that.

I guess. The solution of creating a new protocol for annotating enter_context feels like it works for this specific situation, though, without solving the general problem. Type checkers will continue making bad inferences in other situations regarding the return type of __enter__ for subclasses of ExitStack. If other users hit a simular situation for a user-defined function, they may not know that they need to create a new protocol that ExitStack doesn't inherit from, instead of using the protocol supplied by the standard library.

I really wish we were allowed to do AbstractContextManager[Self] :(

[JelleZijlstra]

Given that AbstractContextManager is a Protocol, it's probably fine to remove the explicit base class from the stub.

I really wish we were allowed to do AbstractContextManager[Self] :(

PEP 673 explicitly disallows this (the example is class Baz(Bar[Self]): ... # Rejected), but I don't think we really discussed it explicitly. It's probably worth bringing up on the python/typing discussions as a possible extension to the PEP.

[sobolevn]

I found this issue while working on python/mypy#13327

With my patch stubtest says:

error: contextlib.AsyncExitStack metaclass missmatch
Stub: at line 151
Missing metaclass
Runtime: at line 537 in file /Users/sobolev/.pyenv/versions/3.8.9/lib/python3.8/contextlib.py
Exiting metaclass: <class 'abc.ABCMeta'>

error: contextlib.ExitStack metaclass missmatch
Stub: at line 136
Missing metaclass
Runtime: at line 468 in file /Users/sobolev/.pyenv/versions/3.8.9/lib/python3.8/contextlib.py
Exiting metaclass: <class 'abc.ABCMeta'>

This happens because AbstractContextManager has ABCMeta metaclass.
Anyway, I guess we can live with that!