Nesting Python context managers
-
13-04-2021 - |
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)
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...