Перехват исключения в менеджере контекста __enter__()

Как это:

import sys

class Context(object):
    def __enter__(self):
        try:
            raise Exception("Oops in __enter__")
        except:
            # Swallow exception if __exit__ returns a True value
            if self.__exit__(*sys.exc_info()):
                pass
            else:
                raise


    def __exit__(self, e_typ, e_val, trcbak):
        print "Now it's running"


with Context():
    pass

Чтобы программа продолжала свою работу без выполнения блока контекста, вам нужно проверить объект контекста внутри блока контекста и выполнить важные действия только в случае успеха __enter__.

class Context(object):
    def __init__(self):
        self.enter_ok = True

    def __enter__(self):
        try:
            raise Exception("Oops in __enter__")
        except:
            if self.__exit__(*sys.exc_info()):
                self.enter_ok = False
            else:
                raise
        return self

    def __exit__(self, e_typ, e_val, trcbak):
        print "Now this runs twice"
        return True


with Context() as c:
    if c.enter_ok:
        print "Only runs if enter succeeded"

print "Execution continues"

Насколько я могу судить, вы не можете полностью пропустить блок with. Обратите внимание, что этот контекст теперь поглощает все исключения. Если вы не хотите проглатывать исключения в случае успеха __enter__, отметьте self.enter_ok в __exit__ и return False, если это True.

Нет. Если есть шанс, что в __enter__() может возникнуть исключение, вам нужно будет перехватить его самостоятельно и вызвать вспомогательную функцию, содержащую код очистки.

Я предлагаю вам следовать RAII (получение ресурсов - это инициализация) и использовать конструктор вашего контекста для потенциально неудачного распределения. Тогда ваш __enter__ может просто вернуть себя, который никогда не должен вызывать исключение. Если ваш конструктор дает сбой, исключение может быть вызвано еще до входа в контекст with.

class Foo:
    def __init__(self):
        print("init")
        raise Exception("booh")

    def __enter__(self):
        print("enter")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("exit")
        return False


with Foo() as f:
    print("within with")

Выход:

init
Traceback (most recent call last):
  File "<input>", line 1, in <module>
  ...
    raise Exception("booh")
Exception: booh

Редактировать: К сожалению, этот подход по-прежнему позволяет пользователю создавать оборванные ресурсы, которые не будут очищены, если он сделает что-то вроде:

foo = Foo() # this allocates resource without a with context.
raise ValueError("bla") # foo.__exit__() will never be called.

Мне очень любопытно, можно ли это обойти, изменив новую реализацию класса или какую-то другую магию Python, которая запрещает создание объектов без контекста with.

Вы можете использовать contextlib.ExitStack (не проверено):

with ExitStack() as stack:
    cm = TstContx()
    stack.push(cm) # ensure __exit__ is called
    with ctx:
         stack.pop_all() # __enter__ succeeded, don't call __exit__ callback

Или пример из документов:

stack = ExitStack()
try:
    x = stack.enter_context(cm)
except Exception:
    # handle __enter__ exception
else:
    with stack:
        # Handle normal case

См. contextlib2 на Python ‹3.3.

если наследование или сложные подпрограммы не требуются, можно использовать более короткий путь:

from contextlib import contextmanager

@contextmanager
def test_cm():
    try:
        # dangerous code
        yield  
    except Exception, err
        pass # do something

class MyContext:
    def __enter__(self):
        try:
            pass
            # exception-raising code
        except Exception as e:
            self.__exit__(e)

    def __exit__(self, *args):
        # clean up code ...
        if args[0]:
            raise

Я сделал это так. Он вызывает __exit__() с ошибкой в ​​качестве аргумента. Если args[0] содержит ошибку, он повторно вызывает исключение после выполнения кода очистки.

Документация содержит пример, который использует contextlib.ExitStack для обеспечения очистки:

Как указано в документации ExitStack.push(), этот метод может быть полезно при очистке уже выделенного ресурса, если последующие шаги в __enter__() реализация не удалась.

Таким образом, вы должны использовать ExitStack() в качестве оборачивающего контекстного менеджера вокруг контекстного менеджера TstContx():

from contextlib import ExitStack

with ExitStack() as stack:
    ctx = TstContx()
    stack.push(ctx)  # Leaving `stack` now ensures that `ctx.__exit__` gets called.
    with ctx:
        stack.pop_all()  # Since `ctx.__enter__` didn't raise it can handle the cleanup itself.
        ...  # Here goes the body of the actual context manager.