嵌套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)

Answer 1

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)

我希望这个能成为标准库的一部分...

Answer 2

创建上下文管理器的简单方法是使用 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 块结束时自动关闭，确保不会出现文件未关闭的情况。