Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Nesting Python context managers

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)
like image 415
Neil G Avatar asked Jan 03 '12 23:01

Neil G


People also ask

What are Python context managers?

Context managers allow you to allocate and release resources precisely when you want to. The most widely used example of context managers is the with statement. Suppose you have two related operations which you'd like to execute as a pair, with a block of code in between.

What does __ enter __ do in Python?

__enter__ and [__exit__] both are methods that are invoked on entry to and exit from the body of "the with statement" (PEP 343) and implementation of both is called context manager. the with statement is intend to hiding flow control of try finally clause and make the code inscrutable.

What is ExitStack?

An ExitStack is (as the name suggests) a stack of clean-up functions. Adding a callback to the stack is the equivalent of calling Go's defer statement. However, clean-up functions are not executed when the function returns, but when execution leaves the with block - and until then, the stack can also be emptied again.

What is Contextlib?

The contextlib module of Python's standard library provides utilities for resource allocation to the with statement. The with statement in Python is used for resource management and exception handling. Therefore, it serves as a good Context Manager.


2 Answers

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

like image 91
Petr Viktorin Avatar answered Sep 19 '22 21:09

Petr Viktorin


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

like image 25
Conchylicultor Avatar answered Sep 22 '22 21:09

Conchylicultor