From a1f09b7f86c0d20941562f5b1254f9c0a2b9a535 Mon Sep 17 00:00:00 2001
From: Michael Schwarz <michi.schwarz@gmail.com>
Date: Mon, 14 Sep 2015 15:03:36 +0200
Subject: Abstraction which allows a running process to be used as a context
 manager.

---
 support/lib/util.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

(limited to 'support')

diff --git a/support/lib/util.py b/support/lib/util.py
index d14b342..48bfd46 100644
--- a/support/lib/util.py
+++ b/support/lib/util.py
@@ -56,7 +56,8 @@ def TemporaryDirectory():
 		shutil.rmtree(dir)
 
 
-def command(args, remove_env = [], set_env = { }, working_dir = None):
+@contextlib.contextmanager
+def command_context(args, remove_env = [], set_env = { }, working_dir = None, use_stderr = False):
 	env = dict(os.environ)
 	
 	for i in remove_env:
@@ -71,10 +72,24 @@ def command(args, remove_env = [], set_env = { }, working_dir = None):
 	except OSError as e:
 		raise UserError('Error running {}: {}', args[0], e)
 	
+	try:
+		yield process
+	except:
+		process.kill()
+		
+		raise
+	finally:
+		process.wait()
+	
 	if process.returncode:
 		raise UserError('Command failed: {}', ' '.join(args))
 
 
+def command(args, remove_env = [], set_env = { }, working_dir = None):
+	with command_context(args, remove_env, set_env, working_dir) as process:
+		process.wait()
+
+
 def bash_escape_string(string):
 	return "'{}'".format(re.sub("'", "'\"'\"'", string))
 
-- 
cgit