From 5f582e3f5246c8a17871a48efc2a52088f0c01b5 Mon Sep 17 00:00:00 2001 From: Valentin Lorentz Date: Thu, 31 Jul 2014 22:52:16 +0200 Subject: [PATCH] Make AtomicFile support context managers. --- src/utils/file.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/utils/file.py b/src/utils/file.py index b2a698033..6c3f9f326 100644 --- a/src/utils/file.py +++ b/src/utils/file.py @@ -156,13 +156,19 @@ class AtomicFile(object): # self.__parent = super(AtomicFile, self) self._fd = open(self.tempFilename, mode) + def __enter__(self): + return self + def __exit__(self, exc_type, exc_value, traceback): + if exc_type: + self.rollback() + else: + self.close() + + @property def closed(self): return self._fd.closed - def close(self): - return self._fd.close() - def write(self, data): return self._fd.write(data)