Context manager protocol extension

A feature that the current context manager protocol lacks and that some people have asked for is the ability for the context manager to skip executionn of the block. Maybe by using a special exception akin to StopIteration the new protocol could support this?

1 Like

After remembering how generator’s send/throw work and playing around with implementation options I really like this approach. I wrote a pure-Python gist which simulates the functioning of this approach.

You were right, getting the rebinding value from StopIteration was really quite easy, but I ended up not liking the “return value of generator is a rebinding value” approach because just forgetting to add return enter_result at the end of the function will rebind with target with None.
For example naive empty cm.__with__ is

def __with__(self):
    yield self

will rebind to None, and the correct way would be

def __with__(self):
    yield self
    return self

So, the new approach as a whole goes like this:
Context manager protocol extends with a new dunder method def __with__(depth: int)
The depth argument is equal to the number of calls of with statement with self as statement expression up the frame stack.

The method must return a generator-iterator when called. This iterator must yield from one to two values. The first value yielded will be bound to the targets in the with statement’s as clause, if any. At the point where the generator yields, the block nested in the with statement is executed. The generator is then resumed after the block is exited. If an unhandled exception occurs in the block, it is reraised inside the generator at the point where the yield occurred. If the exception is catched, the generator must reraise that exception otherwise it is supressed. If an exception has not occurred, or has been suppressed, the generator can yield a second value, which again will be bound to the targets in the with statement’s as clause (if any). After that generator will be closed.


That way all parts of the protocol are explicit. I.e. if you have not explicitly excepted, the exception is reraised. And f you have not explicitly yielded the second value, there is no rebinding. Enter result is available to the exitter because it is closed to the same function.

After more consideration I agree. I think if we will decide to add that rebinding feature, it is better to go with "__leave__ should return a 2-tuple, where first element is suppress flag and second is rebinding value, otherwise it should return bool meaning only “suppress flag”, so if one want to update __exit__ they could just replace it with new signature (with same name for exception) and function behaves like before. However, that way I don’t that much like a rebinding feature, there are workarounds.


I am -1 for that. I don’t want ever to protect myself with

executed = False
with cm:
    executed = True
    ...

if not executed:
    panic

The way I think it should work is

class Manager:
    def can_with(self):
        return ...

    def __enter__(self):
        # In case you as a manager developer want to protect against unsafe execution
        if not self.can_with():
            raise Exception("You should check if manager.can_with() before attempt")
       ...

manager = Manager()
if manager.can_with():
     with manager:
        ...

I see no need for such a feature outside of dirty/hacky one-day solutions.
But with __with__ approach I can see how it could be implemented - just return before first yield. I will reconsicer this if there is a real use case and demand.

2 Likes

This more radical redesign is the kind of thing I was imagining when I heard about “redesigning the context manager”. Seeing your proposal, It looks like it addresses almost all of the points raised very neatly. It would be a little more work if one wants to switch existing code to the new method, but I like this way of doing this.

One thing, in this discussion (about adding a __leave__ method, that PEP 707 linked to), it talks about supporting the case where a base class switches to __leave__ and a derived class still uses __exit__. Would something similar need to be supported with this proposal, where a base class changes to __with__ and the derived class still uses __enter__/__exit__? Or do you just say the whole class hierarchy must use one or the other?


If we look at the original goals from PEP 707. They are:

  1. Better performance in the interpreter. If I understand right, the internal implementation has switched from using the 3-value form of exceptions to the 1-value form, and there’s extra overhead wherever the 3-value forms are still used. If we keep using __exit__, you don’t get the performance benefit of the 1-value exceptions, and likely extra overhead from the interpreter needing to support both forms.

    The “depth” parameter would also add overhead to context managers just to keep track of that information. That goes if they’re used or not, and whether you use __enter__/__exit__ or __enter_ex__/__exit_ex__. I don’t know if this would be significant, or if the savings from using the 1-value exception balances it out.

  2. Simplifying the language. If I understand right, everything that uses the 3-value exceptions should get an alternative with 1-value exceptions, and that that’s been achieved, except for __exit__. The fact that the traceback parameter is rarely used, and it can be gotten through other means, is one of the reasons for it to be removed. Adding a “depth” parameter makes it more complicated again.


When we’re writing a __with__ method. What is the advantage of having the depth parameter passing in, versus doing something like this:

def __with__(self):
    self._depth += 1
    depth = self._depth
    try:
        ... # do setup
        try:
            yield self
        finally:
            ... # do cleanup
    finally:
        self._depth -= 1

It seems easy enough to keep track of the information that way, if you need it. If you don’t, then you leave it out and there’s no performance overhead. Is there a situation where this doesn’t work?

2 Likes

The yield is where the block of code executes, so the yielded value (the first self) is what gets bound. Anything returned at the end is probably going to be ignored. If you want to suppress exceptions, wrap the yield in try/except.

I suggest referencing contextlib.contextmanager, since we already have an implementation of this. All we’d be doing is promoting it to “native” - there shouldn’t be any need to change any other semantics.

Oof, this could be a challenge, yeah. We might have to invert the logic to prefer __enter__/__exit__ if they’re defined (and not-None, so that a subclass can hide them if it wants) for a deprecation period (with a warning), and eventually switch to preferring __with__ and ignoring the old logic if it’s there.[1]

Generally you wouldn’t override __enter__/__exit__ anyway (if it’s designed well, you’d give people specific methods to override), so I imagine most types will be able to get the benefits immediately, even if we keep preferring the old way for now.


  1. A more complex approach could check the MRO and use the most derived implementation’s API. No idea whether that’s worth it though. ↩︎

Look back at this post and this post from the original discussion. Both suggest adding some way to return a value from the context manager when the with block is finished, not just when it starts. That’s where the thing with return values is coming from.

Okay, so the proposal is to enable something like this:

with timeit() as duration:
    code_under_test()

assert isinstance(duration, float)

Which is not possible today because the as name is intended for use within the block and the final primitive value cannot be known until after the block has been executed.

I’d worry about changing the user-visible semantics of with, as opposed to only the implementer’s semantics, though it doesn’t actually seem terrible to simply rebind the name again at the end. And if we go ahead with a generator-based __with__ approach then it’s certainly possible to handle a returned value.

Worth writing it up, at least, but I’d hold onto it loosely. It’s the kind of change that will get a proposal rejected while everything else is uncontroversial. It’s also likely that we’ll find some reason it’s not a good idea in the process of specifying the behaviour.

Here is a concrete example of where these changes will be used. This is not a fictional example, it may well occur in the wild, although perhaps real cases will only use some of the bits we are talking about.
So, it is reentrant context manager which also wants to leave something after the block to represent the result of interaction with the connection.

# What we have to do now.
class ConnectTo:
    def __init__(self, where):
        self.where = where
        self._depth = 0
        self._connection = None
        self._connection_result = None

    @property
    def last_connection_result(self):
        return self._connection_result

    def __enter__(self):
        if self._depth == 0:
            self._connection = get_connection(self.where)

        self._depth += 1
        return self._connection

    def __exit__(self, exc_type, exc, exc_tb):
        if isinstance(exc, ConnectionError):
            self._connection.do_fancy_stuff(exc)
            result = True
        else:
            result = False

        if exc is None:
            self._connection.commit()
            self._connection_result = self._connection.get_state()

        self._depth -= 1
        if self._depth == 0:
            self._connection.disconnect()
            self._connection = None

        return result

with (cm := ConnectTo('db')) as connection:
    ...
result = cm.last_connection_result

# What we can do with __enter_ex__ and __leave__.
class ConnectToEx:
    def __init__(self, where):
        self.where = where
        self._connection = None

    def __enter_ex__(self, depth):
        if depth == 0:
            self._connection = get_connection(self.where)

        return (self._connection, None)

    def __leave__(self, connection, exc, depth):
        result = None
        if isinstance(exc, ConnectionError):
            self._connection.do_fancy_stuff(exc)

        elif exc is None:
            connection.commit()
            result = connection.get_state()

        if depth == 0:
            connection.disconnect()
            self._connection = None

        if exc and not isinstance(exc, ConnectionError):
            raise exc

        return (None, result)

with ConnectToEx('db') as (connection, result):
    ...
result

# What we can do with __with__.
class ConnectToWith:
    def __init__(self, where):
        self.where = where
        self._connection = None

    def __with__(self, depth):
        if depth == 0:
            self._connection = get_connection(self.where)

        try:
            yield (self._connection, None)
        except ConnectionError as e:
            self._connection.do_fancy_stuff(e)
        else:
            self._connection.commit()
            yield (None, self._connection.get_state())

        finally:
            if depth == 0:
                self._connection.disconnect()
                self._connection = None

with ConnectToWith('db') as (connection, result):
    ...
result

As can be seen from this example, everything discussed here can be achieved now, but requires some boilerplate.

The main advantage is that you don’t have to maintain this yourself. It will be quite cheap for the interpreter to do this (interpreter.with_stack.append(id(manager_object)) and interpreter.with_stack.count(id(manager_object))). I’ve faced bugs related to reentrancy several times, and it’s always incorrect handling of the depth counter (e.g. not reducing it in one of the branches). If the counter is inside the interpreter, these bugs should disappear.


I don’t think that’s a problem. If we call this kind of default __with__. In case __enter__/__exit__ code uses super() then it will be an error, and if it doesn’t, as I understand it, it means that the derived class doesn’t want to use the parent’s protocol.

def default___with__(obj, depth):
    # Exception in __enter__ should be propagated.
    enter_result = obj.__enter__()
    try:
        yield enter_result
    except BaseException as e:
        if not obj.__exit__(type(e), e, e.__traceback__):
            raise
    else:
        obj.__exit__(None, None, None)

But as you correctly said, it’s pretty rare when the manager’s creator expects an derived class to overload magic methods, they will make additional methods.
So, the main ‘drawback’ of __with__ approach is that super().__with__ is pretty useless and I expect people will add such methods anyway

class Parent:
    def enter_context(self):
        ...
    def exit_context(self, exc):
        # default handlers

    def __with__(self):
        try:
            yield self.enter_context()
        except Exception as e:
            self.exit_context(e)

class Child(Parent):
    def __with__(self):
        try:
            yield self.enter_context()
        except MyException:
            ...
        except Exception as e:    # this is pretty super().__with__
            self.exit_context(e)

For me, both ideas of depth and rebind are not very important. There are workarounds to get this. It is quite possible that we will go without them, and they will be a potential __with_ex__ if needed later on.

I haven’t worked this through based on your example code but I can imagine simpler ways to manage the boiler-plate rather than changing the context manager protocol such as:

  1. Make a ReentrantContextManager that keeps track of depth but can be subclassed like:
class ConnectToSubclass(ReentrantContextManager):
    # __enter__ and __exit__ supplied by superclass
    def __enter_ex__(...):
        ...
    def __leave__(...):
        ...
  1. Make a decorator like @contextmanager_reentrant that wraps something like the __with__ function or perhaps the class:
@contextmanager_reentrant
class ConnectToDecorator:
    # __enter__ and __exit__ supplied by decorator
    def __init__(...):
        ...
    def __with__(...):
        ...

Maybe there’s also a nice way to do this with generator functions like the existing @contextmanager decorator.

Either way the result could be something that achieves the suggested behaviour but wrapped up to work with the existing context manager protocol rather than requiring a new protocol.

The contrast in the given examples between “what we have to do now” and “what we can do with …” does not suggest to me that this alternative version of the context manager protocol would be easy to understand. The fact that some things are handled implicitly by the interpreter reduces the amount of explicit code but there is instead a cognitive burden to understand how the implicit behaviour of the (nested) with statements interact with the code that is visible. This is especially jarring for the implicit rebinding that takes place invisibly after the with statement. Usually a name is visible at the place in the code where it is rebound even if the binding is implicit like when using a decorator.

Both examples need if depth == 0 in at least two places which hints that it might be cleaner to have separate method(s) to be called for the depth == 0 case. Also apparently the only relevant property of depth is whether it is equal to zero so if there were separate methods for that then maybe the depth argument would not even be needed.

1 Like

Reentrancy is pretty easy with a yield-based context manager:

def reentrant(obj):
    try:
        obj.open()
    except AlreadyOpenError:
        yield obj
    else:
        try:
            yield obj
        finally:
            obj.close()

There are a few important variants on this, which is why we don’t want to bake it into the protocol. Users of it can design the handling they need easily enough, whether it’s using instance state or local state.

3 Likes

My view is that __enter__ and __exit__ are examples of methods that have an “augmenting pattern”; they should always call super because you don’t know what your superclass is, and you don’t know if it will have some important behavior.

I realize that we’re just hashing things out, but in the documentation, I think the examples should always show delegation to super as a good habit to get into. Here’s an example of unconditionally calling super using the AbstractContextManager as a stub provider:

from contextlib import AbstractContextManager
from time import perf_counter_ns
from types import TracebackType
from typing_extensions import Self, override

class Timer(AbstractContextManager['Timer']):
    @override
    def __init__(self) -> None:
        super().__init__()
        self.start = 0
        self.end = 0

    @override
    def __enter__(self) -> Self:
        super().__enter__()
        self.start = perf_counter_ns()
        return self

    @override
    def __exit__(self,
                 exc_type: None | type[BaseException],
                 exc_val: None | BaseException,
                 exc_tb: None | TracebackType,
                 /) -> None:
        super().__exit__(exc_type, exc_val, exc_tb)
        self.end = perf_counter_ns()

I think you’re absolutely right. As far as I can see, there’s no good way to implement inheritance with a __with__ as we see it. The tempting thing to do is to yield from the parent class’s __with__, but the problem with that is that we’re stuck yielding whatever the parent wanted to yield. And we can’t alter that in the local scope or else the parent won’t receive exceptions. Am I missing something?

If this is the case, then I agree with Andrej that this definition of __with__ is essentially broken with respect to inheritance. Adding auxilliary methods is not a reasonable compromise since parent classes (which are unknown) may not know about your auxilliary methods.

Maybe map could be altered to make this work?

    @override
    def __with__(self) -> Generator[Self]:
        self.start = perf_counter_ns()
        try:
            yield from map(lambda _: self, super().__with__())
        finally:
            self.end = perf_counter_ns()

Even so, this doesn’t seem simpler than the original code with __enter__ and __exit__.

That’s not how inheritance (as a design pattern) works. This is encapsulation using subclassing. And Python doesn’t really do either that well - composition is preferred.

For inheritance, the base class provides the entry point to the common functionality and the subclass provides the implementation. In this case, that means the base class must be a context manager. Otherwise, you’re wrapping a context manager around a class that isn’t one, and so you can do whatever you like.

In essence, a base class that doesn’t provide subclassable behaviour isn’t really a base class.

I guess if you really wanted to subclass something that doesn’t provide proper extension points and override its behaviour while keeping the underlying functionality, you could do it this way:

class SubClass:
    def __with__(self):
        with super() as s:
            # do something else
            yield s

I’m pretty sure that would require special handling of super() objects by the with bytecodes, but that’s doable. It does, however, need to be specified.

1 Like

Of course you’re right that composition is preferred (“composition over inheritance”). But you can’t always compose things. Sometimes you need inheritance because of polymorphism requirements. Personally, I think Python’s inheritance pattern is just fine.

Yes, and in the example I gave, the base class is AbstractContextManager, which is a context manager? Am I missing your point?

That’s a fascinating solution! It also gets rid of the try/finally right? If this really can be done, then __with__ looks much more attractive!

Right now it can’t, but presumably this just means we need to implement it:

Python 3.11.3 (tags/v3.11.3:f3909b8, Apr  4 2023, 23:49:59) [MSC v.1934 64 bit (AMD64)] on win32
Type "help", "copyright", "credits" or "license" for more information.
>>> class Base:
...   def __enter__(self):
...     print("Base.__enter__")
...     return 'base'
...   def __exit__(self, *a):
...     print("Base.__exit__")
...
>>> class Sub(Base):
...   def __enter__(self):
...     print("Sub.__enter__")
...     return 'sub'
...   def __exit__(self, *a):
...     print("Sub.__exit__")
...   def f(self):
...     with super() as s:
...       print("Got", s)
...
>>> Sub().f()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 8, in f
TypeError: 'super' object does not support the context manager protocol

FWIW, my expected/hopeful output here would be:

>>> Sub().f()
Base.__enter__
Got base
Base.__exit__
2 Likes

One thing we all should keep in mind while discussing things is that there is also async world and their
async with which should be in line with synchronous world. The main thing is that we cannot suggest any pattern using yield from because this is a syntax error in an asynchronous generator function.

2 Likes

Yes, I’m aware of the ways to make reentrancy. It never was a question of “the language does not allow to make it”, rather a question “what way is the best for that concrete context manager”. And pretty often contextlib and/or some pattern satisfy the need. I wrote about `depth’ to make it heard, but I don’t hold on to that idea in any way. So, I step down with it. If someone has an example where it significantly simplifies the code, they can advocate for it.


But I will still advocate for the “rebind on block exit” idea.
First thing, we already have something simillar

try:
    1 / 0
except Exception as e:
    pass
e  # UnboundLocalError: local variable 'e' referenced before assignment

So, “rebinding the name at exit of some syntax block” is not completly new for the language.
And because it needs an explicit statement in __with__ approach (second yield) I think it does not add that much cognitive burden.
After all, I expect people will use it mostly as:

with cm() as result:
     assert result is None
assert result is not None

with cm() as (enter, leave):
    assert enter is not None and leave is None
assert leave is not None

Am I understanding correctly that in the matter of inheritance our final goal with __with__ approach is


class Base:
    def __with__(self):
        print("Base.__with__ enter")
        try:
            yield self
        except Exception as e:
            print("Base.__with__ exception")
            raise e
        else:
            print("Base.__with__ exit")
        finally:
            print("Base.__with__ finally")


class Derived(Base):
    def __with__(self):
        try:
            with super() as enter_result:
                print("Derived.__with__ enter")
                yield self
        except Exception as e:
            print("Derived.__with__ exception")
            raise e
        else:
            print("Derived.__with__ exit")
        finally:
            print("Derived.__with__ finally")


def exception_case():
    with Derived() as d:
        print("exception_case body")
        raise Exception

    # prints:
    # Base.__with__ enter
    # Derived.__with__ enter
    # exception_case body
    # Base.__with__ exception
    # Base.__with__ finally
    # Derived.__with__ exception
    # Derived.__with__ finally

def no_exception_case():
    with Derived() as d:
        print("no_exception_case body")

    # prints:
    # Base.__with__ enter
    # Derived.__with__ enter
    # no_exception_case body
    # Derived.__with__ exit
    # Derived.__with__ finally
    # Base.__with__ exit
    # Base.__with__ finally
1 Like

I hope so. Except, I think Base should inherit from AbstractContextManager and should also wrap its yield in a with super().

There’s no reason to do this. I don’t even know why we have that, apart from perhaps as documentation for people who read source code instead of docs.

Context managers are a protocol, not an override. Abstract base classes are way too limited and way too heavyweight to normalise them when they aren’t useful.

This is only because of a reference cycle created by the traceback that leads to massive memory leaks when you handle exceptions in a loop. The unbound exception name is a wart, and not a good precedent for future decisions - we don’t want more of these, we want less.

If you want to avoid leaking resources beyond a with statement, then you can clean them up. If except blocks had been implemented after context managers, they’d probably use similar semantics to that, but they came first and we decided to make things better on the second attempt.

So let’s not consider turning warts into principles of good language design :wink:

Okay, it seems before answering the question “how to express rebinding after block” we first need to answer “do we want the protocol allow rebinding at all”. This is a pretty important question, because unlike depth thing it is something the language does not allow to do in a clear way.

To give context to new readers, the problem was first raised here

In code it looks like

from time import perf_counter
class TimeIt:
    def __init__(self):
        self.elapsed = None

    def __enter__(self):
        self.start_time = perf_counter()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.elapsed = perf_counter() - self.start_time

with TimeIt() as time_it:
    ...
print(time_it.elapsed)

# In case manager.__enter__ returns non-self
with (use_me_to_get_result := cm()) as enter:
     ...  # work with 'enter'
use_me_to_get_result.result

The proposed solution is to allow to return extra value (return value of __leave__ or second yielded value of __with__). When such value is present as target of with statement rebinds with that value.

I’d like to hear input from all the people, whenever this “rebind after block” feature is desirable. Maybe things like timeit and assertRaises are black sheeps, rare cases that should use workarounds.

The reason to do this is so that anyone can inherit from multiple context manager classes without losing behavior. If you try to explicitly call parent classes, you could get into trouble when there are common base classes. I understand if people don’t ever want to use inheritance with the context managers that they write. But by not calling super, we prevent anyone else from inheriting from our context managers.

Okay, I see your point about “normalizing” this. People are free to write context managers that don’t support inheritance, just like they’re free not to call super in __init__. I personally like to keep my options open, so I make an effort to call super in both __init__ and __enter__/__exit__.

As long as calling super is easy, then I’m happy.

What’s the “wart” here?