Question

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)
Was it helpful?

Solution

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

OTHER TIPS

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, filename):
    self.filename = filename

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

Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow
scroll top