[PATCH 1/2] Make atomic_ofstream a Context Manager

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 <[email protected]>
---
 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