Python上下文管理器的嵌套问题

Question

在这个问题中，我定义了一个包含上下文管理器的上下文管理器。最容易正确地完成这种嵌套的方式是什么？我最终在self.__enter__()中调用了self.temporary_file.__enter__()。然而，在self.__exit__中，如果出现异常，我非常确定必须在finally块中调用self.temporary_file.__exit__(type_, value, traceback)。如果在self.__exit__中出现问题，我应该设置type_、value和traceback参数吗？我检查了contextlib，但找不到任何有助于解决此问题的实用程序。

原始代码来自问题：

import itertools as it
import tempfile

class WriteOnChangeFile:
    def __init__(self, filename):
        self.filename = filename

    def __enter__(self):
        self.temporary_file = tempfile.TemporaryFile('r+')
        self.f = self.temporary_file.__enter__()
        return self.f

    def __exit__(self, type_, value, traceback):
        try:
            try:
                with open(self.filename, 'r') as real_f:
                    self.f.seek(0)
                    overwrite = any(
                        l != real_l
                        for l, real_l in it.zip_longest(self.f, real_f))
            except IOError:
                overwrite = True
            if overwrite:
                with open(self.filename, 'w') as real_f:
                    self.f.seek(0)
                    for l in self.f:
                        real_f.write(l)
        finally:
            self.temporary_file.__exit__(type_, value, traceback)

- Neil G

Answer 1

使用contextlib.contextmanager是创建上下文管理器的简单方法。例如：

@contextlib.contextmanager
def write_on_change_file(filename):
    with tempfile.TemporaryFile('r+') as temporary_file:
        yield temporary_file
        try:
             ... some saving logic that you had in __exit__ ...

然后使用with write_on_change_file(...) as f:。
with语句块的主体将被执行，而不是yield。 如果您想捕获主体中发生的任何异常，请将yield本身放在一个try块中。
临时文件将始终在其with块结束时正确关闭。

- Petr Viktorin

---

这真的很好。我会把问题保持开放一段时间，以防此问题产生其他好的答案。 - Neil G

使用 @contextlib.contextmanager 很方便，但是仍然有一些情况需要使用手动定义了 __enter__ 和 __exit__ 方法的类。你有关于这样做的建议吗？ - Zearin

好的，当更方便时再执行 - 例如当对象需要做更多事情而不仅仅是作为上下文管理器时（尽管在这种情况下，您还应考虑添加@contextlib.contextmanager方法）。 - Petr Viktorin

我的唯一建议是从对称性的角度出发，将 self.temporary_file = tempfile.TemporaryFile('r+') 移动到 init 语句中。 - hum3

Answer 2

contextlib.contextmanager 很适合用于函数，但当我需要将类作为上下文管理器时，我使用以下实用程序：

class ContextManager(metaclass=abc.ABCMeta):
  """Class which can be used as `contextmanager`."""

  def __init__(self):
    self.__cm = None

  @abc.abstractmethod
  @contextlib.contextmanager
  def contextmanager(self):
    raise NotImplementedError('Abstract method')

  def __enter__(self):
    self.__cm = self.contextmanager()
    return self.__cm.__enter__()

  def __exit__(self, exc_type, exc_value, traceback):
    return self.__cm.__exit__(exc_type, exc_value, traceback)

使用@contextlib.contextmanager可以使用生成器语法声明上下文管理器类。这使得嵌套上下文管理器更加自然，无需手动调用__enter__和__exit__。例如:

class MyClass(ContextManager):

  def __init__(self, filename):
    self._filename = filename

  @contextlib.contextmanager
  def contextmanager(self):
    with tempfile.TemporaryFile() as temp_file:
      yield temp_file
      ...  # Post-processing you previously had in __exit__


with MyClass('filename') as x:
  print(x)

我希望这可以成为标准库中的一部分...

- Conchylicultor

---

@NeilG，根据@Zearin的评论，有些情况下需要使用类而不是函数，在这种情况下无法使用contextlib.contextmanager。此工具允许使用类与contextlib.contextmanager生成器语法。这使得嵌套上下文管理器更加自然，无需存储x.__enter__()并手动调用x.__exit__()。 - Conchylicultor

是的，但这并没有回答我的问题，我的问题是关于嵌套上下文管理器的。此外，无论何时实现__enter__和__exit__，您都应该始终调用super，以防您的类在协作继承中被使用。 - Neil G

@NeilG，我更新了代码片段以更好地匹配问题。基本上，这与Petr Viktorin的答案相同，但使用类进行操作。 - Conchylicultor

现在看起来很不错。据我所知，它允许我们直接使用with instance:来替换with instance._hidden_ctxt1, instance._hidden_ctxt2, instance._hidden_ctxt3:，对吗？（与使用现有的stdlib装饰器相反，我们可以编写with instance.as_with():，需要在with子句中进行额外的方法调用） - Dubslow

啊，不太对，这只允许委托给单个内部上下文。不过，作为一个标准库算法，我认为检查 contextmanager() 产生的值是否可迭代并不难，所以完全可以做到。 - Dubslow