(Async)ExitStack can lose __context__ when its __exit__ runs without an active exception

[oremanj]

contextlib.ExitStack tries to make exception handling work like it would with a bunch of nested with statements. This is difficult because the interpreter doesn't provide a mechanism for setting the active exception from Python code (relevant: #89524) so the __context__ of exceptions raised by context manager __exit__ methods winds up incorrect at first and must be fixed up after-the-fact.

Two years ago, the __context__ fixup was too liberal: it would set things as __context__ that the equivalent nested with blocks wouldn't. Then #88760 was filed, and #27089 was committed to fix it. Now the __context__ fixup is too conservative: it sometimes doesn't set things as context that the equivalent nested with blocks would.

The fundamental problem here is that ExitStack can't tell the difference between a __context__ that is None because the interpreter wanted to chain with the active exception but there wasn't one, and a __context__ that is None because no chaining should have occurred (e.g. because the context manager was reraising an exception from some other place in the program). When ExitStack.__exit__ runs with an active exception, it can look for that exception in the __context__ chain. When there isn't an active exception, matching the behavior of nested with blocks is a lot more complicated.

Solutions I could imagine:

• Do Provide a more convenient way to set an exception as "active", from Python code #89524. It looks like it was closed for lack of input from the original author; I'd be happy to pick it up if there's appetite for the feature.
• Put the entirety of ExitStack.__exit__ within a construct like:

try:
    raise Exception("internal exception used for ExitStack __context__ fixup")
except Exception as frame_exc:
    # now there's an active exception and the existing logic will work

• Do nothing, and accept the divergence between ExitStack and nested with blocks.

(We definitely should not just revert #27089; the problems associated with too much __context__ are in practice far worse than the problems associated with too little. Most recently I ran into python-trio/trio#2577 on 3.8.)

Reproducer (chainer.py):

from contextlib import contextmanager, ExitStack
import traceback

class raise_on_exit:
    def __init__(self, ty):
        self._ty = ty

    def __enter__(self):
        pass

    def __exit__(self, *exc):
        raise self._ty()

def example_regular():
    with raise_on_exit(KeyError):
        with raise_on_exit(ValueError):
            pass

def example_stacked():
    with ExitStack() as stack:
        stack.enter_context(raise_on_exit(KeyError))
        stack.enter_context(raise_on_exit(ValueError))

try:
    example_regular()
except Exception:
    traceback.print_exc()

print("--------")

try:
    example_stacked()
except Exception:
    traceback.print_exc()

Output on 3.8.7 (does not have #27089): the ValueError is correctly left in __context__ in both cases

Traceback (most recent call last):
  File "chainer.py", line 17, in example_regular
    pass
  File "chainer.py", line 12, in __exit__
    raise self._ty()
ValueError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "chainer.py", line 25, in <module>
    example_regular()
  File "chainer.py", line 17, in example_regular
    pass
  File "chainer.py", line 12, in __exit__
    raise self._ty()
KeyError
--------
Traceback (most recent call last):
  File "/opt/hrt/hrtpy38/root/usr/lib/python3.8/contextlib.py", line 510, in __exit__
    if cb(*exc_details):
  File "chainer.py", line 12, in __exit__
    raise self._ty()
ValueError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "chainer.py", line 32, in <module>
    example_stacked()
  File "chainer.py", line 22, in example_stacked
    stack.enter_context(raise_on_exit(ValueError))
  File "/usr/lib/python3.8/contextlib.py", line 525, in __exit__
    raise exc_details[1]
  File "/usr/lib/python3.8/contextlib.py", line 510, in __exit__
    if cb(*exc_details):
  File "chainer.py", line 12, in __exit__
    raise self._ty()
KeyError

Output on 3.10.10 (has #27089): the ValueError context is lost when using ExitStack

Traceback (most recent call last):
  File "[...]/chainer.py", line 16, in example_regular
    with raise_on_exit(ValueError):
  File "[...]/chainer.py", line 12, in __exit__
    raise self._ty()
ValueError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "[...]/chainer.py", line 25, in <module>
    example_regular()
  File "[...]/chainer.py", line 15, in example_regular
    with raise_on_exit(KeyError):
  File "[...]/chainer.py", line 12, in __exit__
    raise self._ty()
KeyError
--------
Traceback (most recent call last):
  File "[...]/chainer.py", line 32, in <module>
    example_stacked()
  File "[...]/chainer.py", line 20, in example_stacked
    with ExitStack() as stack:
  File "/usr/lib/python3.10/contextlib.py", line 576, in __exit__
    raise exc_details[1]
  File "/usr/lib/python3.10/contextlib.py", line 561, in __exit__
    if cb(*exc_details):
  File "[...]/chainer.py", line 12, in __exit__
    raise self._ty()
KeyError

[arhadthedev]

@ncoghlan, @1st1 (as contextlib experts)

[zware]

This has just come up again in a Discourse post about losing exceptions from callbacks.

[bharel]

I think it's better to sort the context bug first before implementing exception groups.

Right now it hides everything underneath and the context gets deleted.

Is there any PR that attempted to fix it with avoiding the cycle error?

[oremanj]

I never submitted a PR because I never received feedback about which of the approaches I proposed might be accepted. I am still happy to solve this problem if there is any indication that my effort will not be wasted.

[bharel]

Yeah, I wish to help you but it feels like there's a lot of context to enter (pun intended).
Will have to set aside some time to figure it out thoroughly.