Context manager protocol extension (2nd attempt)

After writing today a bunch of __exit__(self, exc_type, exc_value, traceback) I would like to ressurect previous discussion with a narrower scope.

I propouse to supersede current pair of __enter__/__exit__ with single __with__ method which behaves the same way as function decorated with contextlib.contextmanager. No changes in syntax or semantic of with statement, and old way to implement protocol remains supported (but probably discouraged and/or undocumented). I.e. current

class Suppress:
    def __init__(self, *types: type[Exception]):
        self.types = types

    def __enter__(self):
        return None

    def __exit__(self, exc_type, exc_value, traceback):
        if isinstance(exc_value, self.types):
            return True
        return False

becomes

class Suppress:
    def __init__(self, *types: type[Exception]):
        self.types = types

    def __with__(self):
        try:
            yield
        except self.types:
            pass

This approach has the following advantages:

  1. IMO, looks more pythonic and easier to understand.
  2. Allows easier transformation of existing try/except handlers into context managers.
  3. Allows access to a context target variable within a manager without storing it in a manager object.
  4. Removes one more case of the type/value/traceback triplet from the language.

However adding it to the langauage will have some disadvantages:

  1. It makes the language more complex in general.
  2. Every tool that works with/wraps context managers will have to handle that class can implement the protocol in two different ways.
  3. Users who want to be compatiable with Python versions without that feature, I think, will prefer __enter__/__exit__ approach because __with__ approach require to store result generator object somewhere.

If this idea gets support from the community/core devs, I’d happily settle on implementation details and add it to the language.

11 Likes

Why should we add this if we already have @contextlib.contextmanager? What does the new feature add?

3 Likes

Two things, as I read it:

  1. It avoids the need to import a standard library module, relying on a built-in protocol instead.

  2. Some real-world classes do other things besides providing a context manager. Implementing the context management as a method means that client code can use an instance directly in with statements rather than needing to use a separate named context manager.

As stands, we choose between:

from contextlib import closing

class Device:
    def __init__(self, path):
        ...
    def read(self, count):
        ...
    def close(self):
        ...

with closing(Device('foo/bar/baz')) as d:
    ...

or

from contextlib import contextmanager

class Device:
    def __init__(self, path):
        ...
    def read(self, count):
        ...

@contextmanager
def managed(device):
    try:
        yield device
    finally:
        # Of course, if this just calls a `close` method then
        # we could have just used the first example; this is
        # just to illustrate the more general form.
        ...

with managed(Device('foo/bar/baz')) as d:
    ...

or else

class Device:
    def __init__(self, path):
        ...
    def read(self, count):
        ...
    def __enter__(self):
        return self
    def __exit__(self, exc_type, exc_value, traceback):
        ...

with Device('foo/bar/baz') as d:
    ...

With the new proposal, we could use the approach to writing the cleanup code that contextlib.contextmanager provides (which is simpler than the protocol around __exit__), avoid the import and still use the instance directly:

class Device:
    def __init__(self, path):
        ...
    def read(self, count):
        ...
    def __with__(self):
        try:
            yield self
        finally:
            ...

with Device('foo/bar/baz') as d:
    ...
4 Likes

@contextlib.contextmanager can not be applied to a class. This feature makes it possible to be the same straightforward of how to replace try/except with a context manager with a class as it is with a contextmanager wrapped function.


Because I started my journey into the Python in 2.7 times, I became familliar with type/value/tb trio pretty fast, but nowadays I think people do wonder why __exit__ has such a strange signature if it is just type(value), value, value.__traceback__, and may be not sure if it is guaranteed that values are really, in that form (and actually before bpo-45711 it was not).

4 Likes

It kind of makes me wonder if contextmanager could be changed to work as a decorator to add __with__ like semantics to a class.

class Square:

    @contextmanager
    def with(self):
        ...

I’m not sure how beneficial it is to add it at the base level when the base level has the existing protocol.

So it sounds like making @contextmanager work for classes would achieve the same thing and would require a lot less changes (which will make it a lot more likely to get accepted) - why not go for that?

You could make a class decorator or use a really ugly hack exploiting __set_name__ on a descriptor to inject the necessary methods[1].

But in either case type checkers would not be able to understand that the class can now act a context manager, unless they added special casing for it like dataclasses.


  1. if you insist the decorator magically needs to work differently on a specifically named method ↩︎

Yeah, as David said, there is a way to do the thing without interpreter support but it is pretty ugly. The thing is that for classes we most likely don’t want wrapped class be converted in something else, which @contextmanager does (it returns a wrapper function that returns _GeneratorContextManager), so it will be a function that modify the class itself to add __enter__/__exit__, a la

Code that does a thing
def __with__manager[T](typ: type[T]) -> type[T]:
    generators: dict[int, Generator] = {}

    def __enter__(self):
        generators[id(self)] = gen = typ.__with__(self)
        try:
            result = next(gen)
        except StopIteration:
            raise RuntimeError("generator does not yielded.")

        return result

    def __exit__(self, typ, value, tb):
        try:
            gen = generators.pop(id(self))
        except KeyError:
            raise RuntimeError("__exit__ called before __enter__")

        if typ is None:
            try:
                next(gen)
            except StopIteration:
                return False
            raise RuntimeError("generator didn't stop")
        else:
            try:
                gen.throw(value)
            except StopIteration:
                return True
            except BaseException as e:
                if e is not value:
                    raise
                return False
            raise RuntimeError("generator didn't stop after throw()")

    typ.__enter__ = __enter__
    typ.__exit__ = __exit__
    return typ

and even if we teach type checkers to recognize it as a context manager I proposed this feature mostly because of other reasons.


To add more context, In 3.11 and 3.12, Irit Katriel (and others) did a great job to reduce the amount of exc_info triplets replacing it with single exception instance, which eventually made it possible to store exceptions as single object in the thread state in gh-30122, so actually now WITH_EXCEPT_START opcode does __exit__(type(value), value, value.__traceback__), so here C side is simpler than Python side.

After that she proposed PEP 707, which SC rejected as too magic.
But in the PEP discussion some core devs voiced (e.g. Serhiy) their like for better context manager protocol which I proposed in previos discussion, but it faded out because I think it had too wide scope.

So, this idea is not because there is no way to achieve this now, but more as a way to push that work futher, removing I think the only remaining Python-side exc_info triplet.

You’re basically asking here for the two functions __enter__ and __leave__ to be created from a single function. I wonder whether something like this would be viable:

class Ctx:
    @contextlib.generatorcontext
    def __enter__(self):
        ...
        yield
        ...
    __exit__ = __enter__.exiter

It’s not quite as clean as you were looking for, but it’s close, and requires no magic (the decorator constructs a perfectly normal function which serves as __enter__, and then attaches an attribute to it to provide the corresponding __exit__ method).

3 Likes

A trick for using contextlib.contextmanager with a class is delegating the __enter__ and __exit__ logic to an context manager created with contextlib.contextmanager. I believe I first saw this in some code by @njs.

from contextlib import contextmanager

class Device:
    
    def __init__(self, path):
        self._ctx = None
        ...
    
    def read(self, count):
        ...
        
    @contextmanager
    def _enter_and_exit(self):
        try:
            ...
            yield self
        finally:
            ...

    def __enter__(self):
        if self._ctx is not None: # context manager is not reentrant
            raise RuntimeError("Device cannot be opened a second time.")
        self._ctx = self._enter_and_exit()
        return self._ctx.__enter__()
            
    def __exit__(self, exc_type, exc_value, traceback):
        if self._ctx is None:
            raise RuntimeError("Device cannot be closed without opening it")
        ctx = self._ctx
        self._ctx = None  # make context manager reusable
        return ctx.__exit__(exc_type, exc_value, traceback)

Aside from the approach you describe (which certainly has precedent, in the form of @property), this could be done with a class-level decorator or a metaclass (assuming the name of the “source” method is either hard-coded or else specified as a decorator argument).

Yeah, property was what led me to think that. This is kinda doing the inverse (property uses 1-3 functions to create one object, this is using one function to create two) but that was the inspiration nonetheless.

Well, to make my life eaiser in my own code I use helper function

def smart_exit(func):
    def wrapper(self, typ, value, tb):
        return func(self, value)
    return wrapper

class Ctx:
    def __enter__(self):
        ...
    @smart_exit
    def __exit__(self, exc):
        ...

I like your snippet, I’m going to try it out to see if it’s convenient.

But as I said the idea is not to find a good workaround, but to see if there is a will to let the PEP 707 successor a go. After all, there was only 6% for status-quo a year ago.

I do want to point out explicitly that none of the workarounds that translate __with__ into the current interface are equivalent because they aren’t really reentrant with the same object, which the OP proposal would be. I am not actually sure if that is a relevant property.

4 Likes

I actually quite like the idea of __with__ (and of course __awith__): it seems much easier to get right on the first try than the existing __[a]enter__/__[a]exit__ pairs. I am somewhat concerned that there could be mines in the finalization of the generator that I can’t see, and I would like someone more knowledgeable about that area to chime in before I would be a full +1 on the proposal.

Reading a bit of the linked past discussion around PEP 707, @storchaka brought up there a couple of very good points about current limitations of __enter__/__exit__ that can be addressed by __with__, including currently not being able to return something from __exit__, which uses the return value to determine whether to suppress an exception, and not having access to the return value of __enter__ in __exit__. With __with__ the equivalent of the return value of __enter__ is of course still available, and it would be natural to use the return value as the result of exiting the context. There would need to be new syntax to use it, though. My immediate thought there is to borrow the return annotation syntax:

class some_context:
    def __with__(self):
        a = 1
        # Is there a use case for passing something back into the context?
        #x = yield a
        yield a
        return a + 1  # or a + (x if x is not None else 1)?

with some_context() as a -> b:
    assert a == 1
    try:
        b
    except NameError:
        pass
    else:
        assert False, "using 'b' within the context should raise"
assert a == 1
assert b == 2
# If `__with__` raises, `b` should still be unbound

I don’t believe that -> is the right syntax to use, since it really has no relation to its use as a return annotation, but this at least conveys the idea.

I would be willing to sponsor a PEP on this topic. I’m not certain of its acceptance, but I’d like to see it tried :slight_smile:

13 Likes

It’s at least as good as any other syntax I can think of, honestly. But I’m not sure that this feature would be especially valuable. Returning like this seems like a fairly uncommon case (despite @storchaka 's motivating examples) and writing return in a generator (which in turn means the implementation is catching a StopIteration and grabbing its .args[0], I suppose?) also sounds awkward. The problem (of communicating a “result” from the with block) can also be solved by dependency injection/inversion of control:

>>> from time import time, sleep
>>> from contextlib import contextmanager
>>> @contextmanager
... def timed(callback=None):
...     try:
...         t = time()
...         yield
...     finally:
...         callback(time() - t)
... 
>>> with timed(print):
...     sleep(0.1)
...     sleep(0.2)
...     sleep(0.3)
...     sleep(0.4)
... 
1.0013315677642822
1 Like

Not if you are modifying an existing library to add metrics, which appears to be what the motivating example is doing. API changes would be slow and most likely only implemented by a handful of clients.

I just found this: Yield-based contextmanager for classes

Apparently it would be possible to add this to contextlib with a quite complex implementation, so there is no fundamental need for this being build into the language. Since this is the second time this exact proposal has come up, it’s probably a good idea to add it.

1 Like

You may read the previous discussion this idea already was discussed.
If I understand correctly there are 2 possibilites

  1. Do not change syntax, and return value of __with__ rebinds target value. This has a drawback of that loosy branch will rebind value to None and people mostly didn’t like the idea of rebinding the variable as being too magic.
  2. Change the syntax to somehow has two variables. This has a drawback that syntax and semantic do change, which make a change less likely to be accepted and also this is the first time for Python to have as target not being set right after the statement

And in general the idea that “context manager has a return value” was not welcomed. Beasts like Timer considered special enouth to have it’s own attribute based result handling. And I do think that, after all, rebinding is not a good idea.


I don’t see a use case for x = yield result as there expected to be exactly one yield, and zero or 2 or more is an error.


I would like to keep the scope of this proposal small to increase the chance of being accepted, things like “manager result variable” or “optinal execution of the body” can be added later as futher extention of __with__.

Since the user can always store a in the with block if they need it later, why not store the return value in the same symbol if __with__ succeeds? I imagine this makes for simpler documentation, along the lines of “on successful completion, a contains the return value”.

In any case, users would need to know whether or not the return value was stored (in either a or b), so maybe there is a use case for with ... else.

Edit: Apologies, I didn’t see that this was discussed in the previous thread.