This allows using a "with statement" together with instances of atomic_ofstream. Allowing for more readable, less error prone and shorter code.
Signed-off-by: Florian Schmaus <f...@geekplace.eu> --- lib/portage/util/__init__.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/lib/portage/util/__init__.py b/lib/portage/util/__init__.py index 0412b2b5911f..bedcbcfe6fcc 100644 --- a/lib/portage/util/__init__.py +++ b/lib/portage/util/__init__.py @@ -11,6 +11,7 @@ __all__ = ['apply_permissions', 'apply_recursive_permissions', 'stack_dicts', 'stack_lists', 'unique_array', 'unique_everseen', 'varexpand', 'write_atomic', 'writedict', 'writemsg', 'writemsg_level', 'writemsg_stdout'] +from contextlib import AbstractContextManager from copy import deepcopy import errno import io @@ -1246,7 +1247,7 @@ def apply_secpass_permissions(filename, uid=-1, gid=-1, mode=-1, mask=-1, stat_cached=stat_cached, follow_links=follow_links) return all_applied -class atomic_ofstream(ObjectProxy): +class atomic_ofstream(AbstractContextManager, ObjectProxy): """Write a file atomically via os.rename(). Atomic replacement prevents interprocess interference and prevents corruption of the target file when the write is interrupted (for example, when an 'out of space' @@ -1287,6 +1288,13 @@ class atomic_ofstream(ObjectProxy): encoding=_encodings['fs'], errors='strict'), mode=mode, **kargs)) + def __exit__(self, exc_type, exc_val, exc_tb): + if exc_type is not None: + self.abort() + else: + self.close() + return None + def _get_target(self): return object.__getattribute__(self, '_file') -- 2.26.2