Intermediate changes

1 files changed, 54 insertions, 9 deletions

diff --git a/yt/python/yt/logger.py b/yt/python/yt/logger.py
index c885fa9f524..7b876312b91 100644
--- a/yt/python/yt/logger.py
+++ b/yt/python/yt/logger.py
@@ -5,7 +5,10 @@ try:
 except ImportError:
     yatest_common = None
 
+import functools
 import logging
+import os
+import re
 
 
 def set_log_level_from_config(logger):
@@ -20,20 +23,47 @@ def set_log_level_from_config(logger):
         logger.setLevel(level=logging.__dict__[logger_config.LOG_LEVEL.upper()])
 
 
-logging.getLogger("yt.packages.requests.packages.urllib3").setLevel(logging.WARNING)
+class SimpleColorizedStreamHandler(logging.StreamHandler):
+    C_LCYAN = "\033[96m"
+    C_LBLUE = "\033[94m"
+    C_LGREEN = "\033[92m"
+    C_LYELLOW = "\033[93m"
+    C_LGRAY = "\033[37m"
+    C_BOLD = "\033[1m"
+    C_END = "\033[0m"
 
-LOGGER = logging.getLogger("Yt")
+    KW = C_LBLUE
+    URL = C_LCYAN + C_BOLD
+    PARAM = C_LGRAY
+    YSON_PARAM = C_LYELLOW
 
-LOGGER.propagate = False
+    RE_KW = functools.partial(lambda p, r, m: p.sub(r, m), re.compile(r"(Perform HTTP \S+ request|Response received)"), r"{}\1{}".format(KW, C_END))
+    RE_HTTP = functools.partial(lambda p, r, m: p.sub(r, m), re.compile(r"(https?://\S+)"), r"{}\1{}".format(URL, C_END))
+    RE_JSON = functools.partial(lambda p, r, m: p.sub(r, m), re.compile(r"([\w'-]+): "), r"{}\1{}: ".format(PARAM, C_END))
+    RE_YSON = functools.partial(lambda p, r, m: p.sub(r, m), re.compile(r"\"([^\";]+?)\"="), "\"{}\\1{}\"=".format(YSON_PARAM, C_END))
 
-set_log_level_from_config(LOGGER)
+    ENABLED = os.environ.get("YT_LOG_LEVEL") == "Debug"
 
-if logger_config.LOG_PATH is None:
-    LOGGER.addHandler(logging.StreamHandler())
-else:
-    LOGGER.addHandler(logging.FileHandler(logger_config.LOG_PATH))
+    terminator = '\n'  # py2 compat
+
+    def _colorize(self, msg):
+        msg = self.RE_KW(msg)
+        msg = self.RE_HTTP(msg)
+        msg = self.RE_JSON(msg)
+        msg = self.RE_YSON(msg)
+        return msg
+
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            stream = self.stream
+            if stream.isatty() and record.levelno == logging.DEBUG and self.ENABLED:
+                msg = self._colorize(msg)
+            stream.write(msg + self.terminator)
+            self.flush()
+        except Exception:
+            self.handleError(record)
 
-BASIC_FORMATTER = logging.Formatter(logger_config.LOG_PATTERN)
 
 formatter = None
 
@@ -45,6 +75,21 @@ def set_formatter(new_formatter):
         handler.setFormatter(new_formatter)
 
 
+logging.getLogger("yt.packages.requests.packages.urllib3").setLevel(logging.WARNING)
+
+LOGGER = logging.getLogger("Yt")
+
+LOGGER.propagate = False
+
+set_log_level_from_config(LOGGER)
+
+if logger_config.LOG_PATH is None:
+    LOGGER.addHandler(SimpleColorizedStreamHandler())
+else:
+    LOGGER.addHandler(logging.FileHandler(logger_config.LOG_PATH))
+
+BASIC_FORMATTER = logging.Formatter(logger_config.LOG_PATTERN)
+
 set_formatter(BASIC_FORMATTER)