Wrap all Python main functions with decorator.

This decorator check if a module was called as the main module and catches exceptions.

1 files changed, 22 insertions, 1 deletions

diff --git a/support/lib/util.py b/support/lib/util.py
index e343b73..d14b342 100644
--- a/support/lib/util.py
+++ b/support/lib/util.py
@@ -1,4 +1,4 @@
-import contextlib, subprocess, tempfile, shutil, re, os
+import sys, contextlib, subprocess, tempfile, shutil, re, os, inspect
 
 
 class UserError(Exception):
@@ -6,6 +6,27 @@ class UserError(Exception):
 		super(UserError, self).__init__(message.format(*args))
 
 
+def main(fn):
+	"""Decorator for "main" functions. Decorates a function that should be called when the containing module is run as a script (e.g. via python -m <module>)."""
+	
+	frame = inspect.currentframe().f_back
+	
+	def wrapped_fn(*args, **kwargs):
+		try:
+			fn(*args, **kwargs)
+		except UserError as e:
+			print >> sys.stderr, 'Error:', e
+			sys.exit(1)
+		except KeyboardInterrupt:
+			sys.exit(2)
+	
+	if frame.f_globals['__name__'] == '__main__':
+		wrapped_fn(*sys.argv[1:])
+	
+	# Allow the main function also to be called explicitly
+	return wrapped_fn
+
+
 def rename_atomic(source_path, target_path):
 	"""
 	Move the file at source_path to target_path.