Wrap all Python main functions with decorator.

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

4 files changed, 28 insertions, 31 deletions

diff --git a/support/asymptote/__main__.py b/support/asymptote/__main__.py
index 91902dd..238332f 100644
--- a/support/asymptote/__main__.py
+++ b/support/asymptote/__main__.py
@@ -1,4 +1,4 @@
-import os, sys, shutil
+import os, shutil
 from lib import util
 
 
@@ -6,6 +6,7 @@ def _asymptote(in_path, out_path, asymptote_dir, working_dir):
 	util.command([os.environ['ASYMPTOTE'], '-f', 'pdf', '-o', out_path, in_path], set_env = { 'ASYMPTOTE_DIR': asymptote_dir }, working_dir = working_dir)
 
 
+@util.main
 def main(in_path, out_path):
 	_, out_suffix = os.path.splitext(out_path)
 	
@@ -23,12 +24,3 @@ def main(in_path, out_path):
 			shutil.copyfile(temp_out_path, out_path)
 	else:
 		raise Exception('Unknown file type: {}'.format(out_suffix))
-
-
-try:
-	main(*sys.argv[1:])
-except util.UserError as e:
-	print 'Error:', e
-	sys.exit(1)
-except KeyboardInterrupt:
-	sys.exit(2)
diff --git a/support/inkscape/__main__.py b/support/inkscape/__main__.py
index 03f6689..a2c9341 100644
--- a/support/inkscape/__main__.py
+++ b/support/inkscape/__main__.py
@@ -1,4 +1,4 @@
-import sys, os, shutil
+import os, shutil
 from lib import util
 from . import effect, inkscape
 
@@ -36,6 +36,7 @@ def _unfuck_svg_document(temp_svg_path):
 	command_line.run()
 
 
+@util.main
 def main(in_path, out_path):
 	_, out_suffix = os.path.splitext(out_path)
 	
@@ -56,12 +57,3 @@ def main(in_path, out_path):
 			export_effect.write_asy(file)
 		else:
 			raise Exception('Unknown file type: {}'.format(out_suffix))
-
-
-try:
-	main(*sys.argv[1:])
-except util.UserError as e:
-	print 'Error:', e
-	sys.exit(1)
-except KeyboardInterrupt:
-	sys.exit(2)
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.
diff --git a/support/openscad/__main__.py b/support/openscad/__main__.py
index 28de2cf..0eae27e 100644
--- a/support/openscad/__main__.py
+++ b/support/openscad/__main__.py
@@ -1,4 +1,4 @@
-import os, sys
+import os
 from lib import util
 
 
@@ -10,6 +10,7 @@ def _write_dependencies(path, target, dependencies):
 	util.write_file(path, '{}: {}\n'.format(target, ' '.join(dependencies)).encode())
 
 
+@util.main
 def main(in_path, out_path, deps_path):
 	cwd = os.getcwd()
 	
@@ -43,12 +44,3 @@ def main(in_path, out_path, deps_path):
 		# Write output files.
 		_write_dependencies(deps_path, relpath(out_path), deps - ignored_files)
 		util.rename_atomic(temp_out_path, out_path)
-
-
-try:
-	main(*sys.argv[1:])
-except util.UserError as e:
-	print 'Error:', e
-	sys.exit(1)
-except KeyboardInterrupt:
-	sys.exit(2)