colorized dbt output (#441)

colorize dbt output

dbt/config.py

@@ -25,14 +25,12 @@ def read_profile(profiles_dir):
 
 def read_config(profiles_dir):
     profile = read_profile(profiles_dir)
-    return profile.get('config')
+    return profile.get('config', {})
 
 
-def send_anonymous_usage_stats(profiles_dir):
-    config = read_config(profiles_dir)
+def send_anonymous_usage_stats(config):
+    return config.get('send_anonymous_usage_stats', True)
 
-    if config is not None \
-       and not config.get("send_anonymous_usage_stats", True):
-        return False
 
-    return True
+def colorize_output(config):
+    return config.get('use_colors', True)

dbt/logger.py

@@ -1,17 +1,38 @@
 import dbt.clients.system
+import dbt.compat
 import logging
 import os
 import sys
 
+import colorama
+
 # disable logs from other modules, excepting CRITICAL logs
 logging.getLogger('botocore').setLevel(logging.CRITICAL)
 logging.getLogger('contracts').setLevel(logging.CRITICAL)
 logging.getLogger('requests').setLevel(logging.CRITICAL)
 logging.getLogger('urllib3').setLevel(logging.CRITICAL)
 logging.getLogger('snowflake.connector').setLevel(logging.CRITICAL)
 
+
+# Colorama needs some help on windows because we're using logger.info
+# intead of print(). If the Windows env doesn't have a TERM var set,
+# then we should override the logging stream to use the colorama
+# converter. If the TERM var is set (as with Git Bash), then it's safe
+# to send escape characters and no log handler injection is needed.
+colorama_stdout = sys.stdout
+colorama_wrap = True
+
+if sys.platform == 'win32' and not os.environ.get('TERM'):
+    colorama_wrap = False
+    colorama_stdout = colorama.AnsiToWin32(sys.stdout).stream
+
+elif sys.platform == 'win32':
+    colorama_wrap = False
+
+colorama.init(wrap=colorama_wrap)
+
 # create a global console logger for dbt
-stdout_handler = logging.StreamHandler(sys.stdout)
+stdout_handler = logging.StreamHandler(colorama_stdout)
 stdout_handler.setFormatter(logging.Formatter('%(message)s'))
 stdout_handler.setLevel(logging.INFO)
 
@@ -26,6 +47,16 @@ def make_log_dir_if_missing(log_dir):
     dbt.clients.system.make_directory(log_dir)
 
 
+class ColorFilter(logging.Filter):
+    def filter(self, record):
+        subbed = dbt.compat.to_string(record.msg)
+        for escape_sequence in dbt.ui.colors.COLORS.values():
+            subbed = subbed.replace(escape_sequence, '')
+        record.msg = subbed
+
+        return True
+
+
 def initialize_logger(debug_mode=False, path=None):
     global initialized, logger, stdout_handler
 
@@ -49,6 +80,9 @@ def initialize_logger(debug_mode=False, path=None):
             backupCount=7,
         )
 
+        color_filter = ColorFilter()
+        logdir_handler.addFilter(color_filter)
+
         logdir_handler.setFormatter(
             logging.Formatter('%(asctime)-18s: %(message)s'))
         logdir_handler.setLevel(logging.DEBUG)

dbt/main.py

@@ -19,6 +19,7 @@
 import dbt.tracking
 import dbt.config as config
 import dbt.adapters.cache as adapter_cache
+import dbt.ui.printer
 
 
 def main(args=None):
@@ -39,11 +40,15 @@ def handle(args):
 
     # this needs to happen after args are parsed so we can determine the
     # correct profiles.yml file
-    if not config.send_anonymous_usage_stats(parsed.profiles_dir):
+    profile_config = config.read_config(parsed.profiles_dir)
+    if not config.send_anonymous_usage_stats(profile_config):
         dbt.tracking.do_not_track()
     else:
         dbt.tracking.initialize_tracking()
 
+    if dbt.config.colorize_output(profile_config):
+        dbt.ui.printer.use_colors()
+
     res = run_from_args(parsed)
     dbt.tracking.flush()