Make __exit__ return types of generator context managers specifiable

[johnslavik]

Closes #13512

Please see https://discuss.python.org/t/support-suppressing-generator-context-managers/103615 first!

[johnslavik]

I think we can rename the new typevar to _ExcReturnT because that's basically what it means: on exception, return either None or Never (every return type includes Never).

[johnslavik]

This PR so far has a significant drawback, as it assumes all @contextmanager-derived context decorators can suppress exceptions.

Therefore for

from collections.abc import Generator

@contextmanager
def not_suppressing() -> Generator[None]:
    yield

@not_suppressing()
def foo(x: int, y: int) -> int:
    return x // y

it causes foo() to be typed as (x: int, y: int) -> int | None, which is incorrect, as foo() always returns int.

Ideally, whether suppression happens or not could be inferred by the type checker.

[johnslavik]

jax error

It's because Fn (source) can be bound to a consistent subtype of Callable[P, R] (the type that replaced the _F/_AF type variables in this PR) that other consistent subtypes of Callable[P, R] aren't assignable to, e.g. T = TypeVar("T", bound=Callable[[], object]) can be bound to type[object] and def () -> object: isn't assignable to it. Assignability to Fn on the return type was previously guaranteed by _F.

I believe that in this sense ContextDecorator doesn't have to be truly stable, therefore these annotations may have to change downstream if we accept this PR at some point.

prefect error

silence_docker_warnings incorrectly extended decoratees's return type with None because of this PR's incorrect assumptions.

openlibrary error

Similarly for openlibrary, elapsed_time incorrectly extended decoratee's return type with None because of this PR's incorrect assumptions.

[johnslavik]

This PR is blocked until I figure out a way to make this edge case an opt-in.

That might either be via type checker special casing:

def example() -> int:
    with suppress():
        return 1

for both mypy and pyright, type checking this code errors out (although type checkers could be smarter in this exact case :-)):

from contextlib import suppress

def example() -> int:
    # [pyright] error: Function with declared return type "int" must return value on all code paths
    # [mypy] error: Missing return statement  [return]
    with suppress():
        return 1

and this doesn't:

from contextlib import nullcontext

def example() -> int:
    with nullcontext():
        return 1

So it definitely is possible to infer what _ExcReturnT for ContextDecorator should be, but the info has to be inferred from generators. To make this possible, _GeneratorContextManager and _AsyncGeneratorContextManager could adopt the new type variable to parametrize the ContextDecorator / AsyncContextDecorator base.

@contextmanager
def example() -> Generator[None]:
    # type checker should flag this as suppressing context manager
    # therefore `example` becomes a `Callable[_P, _GeneratorContextManager[..., None]]`
    with suppress():
        yield

@contextmanager
def example() -> Generator[None]:
    # normal path, this becomes a `Callable[_P, _GeneratorContextManager[..., Never]]`
    with nullcontext():
        yield

[johnslavik]

Oh, although that doesn't solve the problem with jax.

That shouldn't really be a huge issue if the annotation can be corrected downstream, question is how.
I think that Callable[P, R] should suffice. It better depicts what happens, because what is happening is the callable gets replaced with a possibly different assignable callable that can have different state or behavior (it does).

[johnslavik]

I exposed _ExcReturnT through _[Async]GeneratorContextManager with a default of Never. This should be fully backwards-compatible (except for the jax error) and reproduce the reported false positive. Now I'll need to research possibilities of getting type checkers to support this.

[johnslavik]

OK, this looks good.

[johnslavik]

This would not solve a related problem though:

from contextlib import contextmanager, suppress
from collections.abc import Generator

@contextmanager
def ignore() -> Generator[None]:
    with suppress(ZeroDivisionError):
        yield

def div(a: int, b: int) -> float:
    with ignore():
        return a / b

That should cause in the same error as in

from contextlib import suppress

def div(a: int, b: int) -> float:
    with suppress(ZeroDivisionError):
        return a / b

Which signalizes that _ExitT_co of Abstract[Async]ContextManager should be inferred from generators as well, but that's a separate problem to solve.

[johnslavik]

This led me to seeing a more general problem consisting of at least these 3 issues with how type checkers (don't) handle generator-based context managers:

1. Inferring _ExitT_co argument to AbstractContextManager in generator-based context managers by checking try-except blocks (that sometimes may only be reachable after recursively solving generator delegation) + same for async

2. Inferring _ExitT_co argument to AbstractContextManager in generator-based context managers by checking CM.__exit__ return type from with CM: yield blocks (that sometimes may only be reachable after recursively solving generator delegation) + same for async

3. Inferring _ExcReturnT argument to ContextDecorator that I added in this PR (possibly through AbstractContextManager) in generator-based context managers + same for async
   the _ExcReturnT is added to decorated callable's return type and is None if the exception can be suppressed by the decorating context manager or Never otherwise to retain the original return type

In my understanding, there are only two correct bindings of _ExitT_co and _ExcReturnT together

• Literal[True] + None
• Literal[False] | None + Never

Their individual upper bounds would remain the same, but there should be some way to enforce that logic.

This sheds new light on this issue as a much more general problem that needs discussion and new sections in the typing specification.

[johnslavik]

Details

pwndbg (https://github.com/pwndbg/pwndbg)
+ pwndbg/dbg/gdb/__init__.py: note: In member "evaluate_expression" of class "GDBFrame":
+ pwndbg/dbg/gdb/__init__.py:141: error: "ABC" has no attribute "__enter__"  [attr-defined]
+ pwndbg/dbg/gdb/__init__.py:141: error: "ABC" has no attribute "__exit__"  [attr-defined]

sphinx (https://github.com/sphinx-doc/sphinx)
+ sphinx/builders/__init__.py: note: In member "build" of class "Builder":
+ sphinx/builders/__init__.py:405:13: error: "ABC" has no attribute "__enter__"  [attr-defined]
+ sphinx/builders/__init__.py:405:13: error: "ABC" has no attribute "__exit__"  [attr-defined]
+ sphinx/builders/__init__.py: note: In member "_write_serial" of class "Builder":
+ sphinx/builders/__init__.py:773:13: error: "ABC" has no attribute "__enter__"  [attr-defined]
+ sphinx/builders/__init__.py:773:13: error: "ABC" has no attribute "__exit__"  [attr-defined]

core (https://github.com/home-assistant/core)
+ homeassistant/components/qnap/coordinator.py:74: error: "ABC" has no attribute "__enter__"  [attr-defined]
+ homeassistant/components/qnap/coordinator.py:74: error: "ABC" has no attribute "__exit__"  [attr-defined]

This is odd. Does mypy think that the common base of nullcontext() and @contextmanager(Callable[..., Iterator[None]]) is ABC?
I can't reproduce.

[johnslavik]

Oh, this used to pass on mypy 1.16.1 but it doesn't on mypy 1.17.1...

typeshed on  contextdecorator-call-ret-type ~ typeshed 
❯ bat t.py
───────┬─────────────────────────────────────────────────────────────────────────────────────────────────
       │ File: t.py
───────┼─────────────────────────────────────────────────────────────────────────────────────────────────
   1   │ import random   
   2   │ from contextlib import contextmanager, nullcontext
   3   │ from collections.abc import Iterator
   4   │ 
   5   │ @contextmanager
   6   │ def foobar() -> Iterator[None]:
   7   │     yield
   8   │ 
   9   │ ctxmgr = foobar() if random.randint(0, 1) else nullcontext()
  10   │ reveal_type(ctxmgr)
───────┴─────────────────────────────────────────────────────────────────────────────────────────────────

typeshed on  contextdecorator-call-ret-type ~ typeshed 
❯ uvx 'mypy==1.16.1' --custom-typeshed-dir . t.py
t.py:10: note: Revealed type is "contextlib.AbstractContextManager[None, Union[builtins.bool, None]]"
Success: no issues found in 1 source file

typeshed on  contextdecorator-call-ret-type ~ typeshed 
❯ uvx 'mypy==1.17.1' --custom-typeshed-dir . t.py
t.py:10: note: Revealed type is "abc.ABC"
Success: no issues found in 1 source file

Looks like I spotted a regression in the meantime...

[johnslavik]

Anyways, I'm pretty confident that this is correct. Ready for review.

[johnslavik]

Looks like I spotted a regression in the meantime...

Perhaps it's because of disjoint bases? I'll investigate soon.

[srittau]

I'm not sure we can really use Never to say "may raise an exception". Every function can potentially raise an exception – for example KeyboardInterrupt. As Never is the bottom type, X | Never is equivalent to just X:

from typing import Never

def foo() -> Never | int:
    return 3

reveal_type(foo())  # Revealed type is "builtins.int"
foo() + 5  # legal

[johnslavik]

X | Never is equivalent to just X

Yes, that's the entire point of this solution!
We don't use Never to signal something may raise an exception, but to not extend the return type with None when no suppression happens. This way we leave the return type unchanged, because we union it with Never.

It's an intentional design that's exactly based on what stems from your explanation of Never :-)

[srittau]

I see, makes sense.

[johnslavik]

I think we can make this comment a little clearer, FWIW. Would you like me to do that? @srittau

[github-actions]

Diff from mypy_primer, showing the effect of this PR on open source code:

pwndbg (https://github.com/pwndbg/pwndbg)
+ pwndbg/dbg/gdb/__init__.py: note: In member "evaluate_expression" of class "GDBFrame":
+ pwndbg/dbg/gdb/__init__.py:141: error: "ABC" has no attribute "__enter__"  [attr-defined]
+ pwndbg/dbg/gdb/__init__.py:141: error: "ABC" has no attribute "__exit__"  [attr-defined]

sphinx (https://github.com/sphinx-doc/sphinx)
+ sphinx/builders/__init__.py: note: In member "build" of class "Builder":
+ sphinx/builders/__init__.py:405:13: error: "ABC" has no attribute "__enter__"  [attr-defined]
+ sphinx/builders/__init__.py:405:13: error: "ABC" has no attribute "__exit__"  [attr-defined]
+ sphinx/builders/__init__.py: note: In member "_write_serial" of class "Builder":
+ sphinx/builders/__init__.py:773:13: error: "ABC" has no attribute "__enter__"  [attr-defined]
+ sphinx/builders/__init__.py:773:13: error: "ABC" has no attribute "__exit__"  [attr-defined]

core (https://github.com/home-assistant/core)
+ homeassistant/components/qnap/coordinator.py:74: error: "ABC" has no attribute "__enter__"  [attr-defined]
+ homeassistant/components/qnap/coordinator.py:74: error: "ABC" has no attribute "__exit__"  [attr-defined]

jax (https://github.com/google/jax)
+ jax/_src/api.py:3094: error: Incompatible return value type (got "Callable[[VarArg(Any), KwArg(Any)], Any]", expected "F")  [return-value]

pytest (https://github.com/pytest-dev/pytest)
+ src/_pytest/logging.py:931: error: "ABC" has no attribute "__enter__"  [attr-defined]
+ src/_pytest/logging.py:931: error: "ABC" has no attribute "__exit__"  [attr-defined]

cki-lib (https://gitlab.com/cki-project/cki-lib)
- tests/test_metrics.py:75: error: Call to untyped function "dummy_function" in typed context  [no-untyped-call]