CODEWITHSUNDEEP
pythonwith-statement
Chen Wei
Chen Wei

Reputation: 402

wrap python with statement from 2 classes in a single class

I have a python class A with with statement and class B with with statement. Now it works like:

with A() as a:
    with B() as b:
        do_things()

How can I build a class C that wrap class A and class B so that I can call it like: 

 with C() as c:
        do_things()

has the identical funtions

Upvotes: 1

Views: 859

Answers (3)

user2357112
user2357112

Reputation: 281683

If your reason for doing this is to reduce indentation or merge the with statements, you don't need to. You can just do

with A() as a, B() as b:
    ...

entering multiple context managers on a single line.

If your C has other reasons to exist, you need to be careful about handling the case where one context manager fails after another context manager has been created and/or entered. contextlib.ExitStack can help you implement this robustly, or contextlib2.ExitStack if you're still on Python 2:

import contextlib

class C(object):
    def __enter__(self):
        with contextlib.ExitStack() as stack:
            stack.enter_context(A())
            stack.enter_context(B())
            self._stack = stack.pop_all()
        return self
    def __exit__(self, exc_type, exc_value, exc_tb):
        return self._stack.__exit__(exc_type, exc_value, exc_tb)

Upvotes: 1

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 477318

Short answer: it is possible. But context-managers allow to implement some logic, which makes it "tricky" to implemented it exactly the right way. Below you see a "Proof of Concept", but I do not guarantee that it has exactly the same behavior. Therefore I really advise to work with nested withs.

What is not covered here: the __init__, or __enter__ can raise exceptions as well, and these are then handled by the "outer" context managers. This makes it of course rather complicated. You basically would need to "build" a stack in the __enter__, and then "pop" the stack in case one of the __enter__s fails. This scenario is not covered here.

We can make a "composite" context manager:

class C:

    def __init__(self, *ctxs):
        self.ctxs = ctxs

    def __enter__(self):
        return tuple(ctx.__enter__() for ctx in self.ctxs)

    def __exit__(self, self, exception_type, exception_value, traceback):
        for ctx in reversed(self.ctxs):
            try:
                if ctx.__exit__(exception_type, exception_value, traceback):
                    (exception_type, exception_value, traceback) = (None,) * 3
            except Exception as e:
                exception_value = e
                traceback = e.__traceback__
                exception_type = type(e)
        return exception_value is None

The __exit__ part is tricky. First of all, we need to exit in reverse order. But the exception handling is even more complicated: if an __exit__ silenced an exception, by returning a "truthful" value, then we should pass (None, None, None) as (exception_type, exeption_value, traceback), but a problem that can occur is that an __exit__ on the other hand triggers an exception itself, and thus then introduces a new exception.

We can then use the context processor like:

with C(A(), B()) as (a,b):
    # ...
    pass

The above thus allows to implement a context manager for an arbitrary number of "sub-contextmanagers". We can subclass this to generate a specific one, like:

class ContextAB(C):

    def __init__(self):
        super(ContextAB, self).__init__(A(), B())

and then use this as:

with ContextAB() as (a, b):
    # ...
    pass

But long story short: use nested with statements. It also makes it more explicit what is going on here. Right now the C encapsulate all sorts of logic, which are better made explicit. If entering B fails, then this should result in an exception that is handled by the __exit__ of A, etc. This makes it very cumbersome to get the "details" completely equivalent to the semantics of the with statement.

Upvotes: 3

cs95
cs95

Reputation: 402922

I'd like to suggest an alternative. You can initialise both a and b on the same line:

with A() as a, B() as b:
    do_things()

This is more concise and reduces the amount of indentation in deeply nested code.

However, if you absolutely must use a class, then override the __enter__ and __exit__ methods:

class C:
    def __enter__(self):
        self._a = A()
        self._b = B()
        return (self._a.__enter__(), self._b.__enter__())

    def __exit__(self ,type, value, traceback):
        # Cleanup code here.
        self._b.__exit__(type, value, traceback)
        self._a.__exit__(type, value, traceback)

And then use C inside a context manager like this:

with C() as (a, b):
    do_things()

If you don't want to reference a and b, or if you don't plan on doing anything with them, then

with C():
    do_things()

Will also work. This should be enough to get started, but please note there are drawbacks as kindly mentioned by users in the comments. The main one being that if self._b.__enter__ throws errors, self._a.__enter__ will need to be cleaned up (this can be done using try-except-finally). Furthermore, some context managers may need to be treated differently based on what resource is being managed.

Upvotes: 3

Related Questions