8
votes

In this question, I defined a context manager that contains a context manager. What is the easiest correct way to accomplish this nesting? I ended up calling self.temporary_file.__enter__() in self.__enter__(). However, in self.__exit__, I am pretty sure I have to call self.temporary_file.__exit__(type_, value, traceback) in a finally block in case an exception is raised. Should I be setting the type_, value, and traceback parameters if something goes wrong in self.__exit__? I checked contextlib, but couldn't find any utilities to help with this.

Original code from question:

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)
2

2 Answers

10
votes

The easy way to create context managers is with contextlib.contextmanager. Something like this:

@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__ ...

Then use with write_on_change_file(...) as f:.
The body of the with statement will be executed “instead of” the yield. Wrap the yield itself in a try block if you want to catch any exceptions that happen in the body.

The temporary file will always be properly closed (when its with block ends).

2
votes

contextlib.contextmanager works great for functions, but when I need a classes as context manager, I'm using the following util:

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)

This allow to declare contextmanager classes with the generator syntax from @contextlib.contextmanager. It makes it much more natural to nest contextmanager, without having to manually call __enter__ and __exit__. Example:

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)

I wish this was in the standard library...