Merge pull request #143 from filbranden/color_when

Support --color option to indicate whether to use colors in the terminal

bootstrapvz/base/log.py

@@ -4,18 +4,20 @@ both to a file and to the console.
 import logging
 
 
-def get_console_handler(debug):
+def get_console_handler(debug, colorize):
 	"""Returns a log handler for the console
 	The handler color codes the different log levels
 
 	:params bool debug: Whether to set the log level to DEBUG (otherwise INFO)
+	:params bool colorize: Whether to colorize console output
 	:return: The console logging handler
 	"""
 	# Create a console log handler
 	import sys
 	console_handler = logging.StreamHandler(sys.stderr)
-	# We want to colorize the output to the console, so we add a formatter
-	console_handler.setFormatter(ConsoleFormatter())
+	if colorize:
+		# We want to colorize the output to the console, so we add a formatter
+		console_handler.setFormatter(ConsoleFormatter())
 	# Set the log level depending on the debug argument
 	if debug:
 		console_handler.setLevel(logging.DEBUG)

bootstrapvz/base/main.py

@@ -34,7 +34,7 @@ def main():
 def get_opts():
 	"""Creates an argument parser and returns the arguments it has parsed
 	"""
-	from docopt import docopt
+	import docopt
 	usage = """bootstrap-vz
 
 Usage: bootstrap-vz [options] MANIFEST
@@ -44,10 +44,15 @@ Options:
                      If <path> is `-' file logging will be disabled.
   --pause-on-error   Pause on error, before rollback
   --dry-run          Don't actually run the tasks
+  --color=auto|always|never
+                     Colorize the console output [default: auto]
   --debug            Print debugging information
   -h, --help         show this help
 	"""
-	return docopt(usage)
+	opts = docopt.docopt(usage)
+	if opts['--color'] not in ('auto', 'always', 'never'):
+		raise docopt.DocoptExit('Value of --color must be one of auto, always or never.')
+	return opts
 
 
 def setup_loggers(opts):
@@ -68,7 +73,15 @@ def setup_loggers(opts):
 		file_handler = log.get_file_handler(path=logpath, debug=True)
 		root.addHandler(file_handler)
 
-	console_handler = log.get_console_handler(debug=opts['--debug'])
+	if opts['--color'] == 'never':
+		colorize = False
+	elif opts['--color'] == 'always':
+		colorize = True
+	else:
+		# If --color=auto (default), decide whether to colorize by whether stderr is a tty.
+		import os
+		colorize = os.isatty(2)
+	console_handler = log.get_console_handler(debug=opts['--debug'], colorize=colorize)
 	root.addHandler(console_handler)