From 9fdeea7bd444c8239dad008a67306fc28542a8c2 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <ruslan@iterative.ai>
Date: Sun, 9 Feb 2020 20:31:23 +0200
Subject: [PATCH] logger: allow exc_info in all record types

This better reflects the original logger formatter.

This also fixes annoying redundant newline between traceback and "need
any help" footer.
---
 dvc/logger.py             | 29 ++++++++++-------------------
 tests/unit/test_logger.py | 16 +++++++---------
 2 files changed, 17 insertions(+), 28 deletions(-)

diff --git a/dvc/logger.py b/dvc/logger.py
index 9385a6aa2a..d1cf08e83a 100644
--- a/dvc/logger.py
+++ b/dvc/logger.py
@@ -51,7 +51,6 @@ class ColorFormatter(logging.Formatter):
 
     color_code = {
         "DEBUG": colorama.Fore.BLUE,
-        "INFO": "",
         "WARNING": colorama.Fore.YELLOW,
         "ERROR": colorama.Fore.RED,
         "CRITICAL": colorama.Fore.RED,
@@ -59,29 +58,21 @@ class ColorFormatter(logging.Formatter):
 
     def format(self, record):
         msg = record.msg.format(*record.args) if record.args else record.msg
+        exception, stack_trace = self._parse_exc(record)
+        return ("{prefix}{description}{stack_trace}").format(
+            prefix=self._prefix(record),
+            description=self._description(msg, exception),
+            stack_trace=stack_trace,
+        )
 
+    def _prefix(self, record):
         if record.levelname == "INFO":
-            return msg
-
-        if record.levelname == "ERROR" or record.levelname == "CRITICAL":
-            exception, stack_trace = self._parse_exc(record)
-
-            return (
-                "{color}{levelname}{nc}: {description}" "{stack_trace}\n"
-            ).format(
-                color=self.color_code.get(record.levelname, ""),
-                nc=colorama.Fore.RESET,
-                levelname=record.levelname,
-                description=self._description(msg, exception),
-                msg=msg,
-                stack_trace=stack_trace,
-            )
+            return ""
 
-        return "{color}{levelname}{nc}: {msg}".format(
+        return "{color}{levelname}{nc}: ".format(
             color=self.color_code.get(record.levelname, ""),
-            nc=colorama.Fore.RESET,
             levelname=record.levelname,
-            msg=msg,
+            nc=colorama.Fore.RESET,
         )
 
     def _current_level(self):
diff --git a/tests/unit/test_logger.py b/tests/unit/test_logger.py
index 0847a5e451..1b7d906e11 100644
--- a/tests/unit/test_logger.py
+++ b/tests/unit/test_logger.py
@@ -44,7 +44,7 @@ def test_error(self, caplog):
         with caplog.at_level(logging.INFO, logger="dvc"):
             logger.error("message")
 
-            expected = "{red}ERROR{nc}: message\n".format(**colors)
+            expected = "{red}ERROR{nc}: message".format(**colors)
 
             assert expected == formatter.format(caplog.records[0])
 
@@ -55,7 +55,7 @@ def test_exception(self, caplog):
             except Exception:
                 logger.exception("message")
 
-            expected = "{red}ERROR{nc}: message\n".format(**colors)
+            expected = "{red}ERROR{nc}: message".format(**colors)
 
             assert expected == formatter.format(caplog.records[0])
 
@@ -66,7 +66,7 @@ def test_exception_with_description_and_without_message(self, caplog):
             except Exception:
                 logger.exception("")
 
-            expected = "{red}ERROR{nc}: description\n".format(**colors)
+            expected = "{red}ERROR{nc}: description".format(**colors)
 
             assert expected == formatter.format(caplog.records[0])
 
@@ -77,9 +77,7 @@ def test_exception_with_description_and_message(self, caplog):
             except Exception:
                 logger.exception("message")
 
-            expected = "{red}ERROR{nc}: message - description\n".format(
-                **colors
-            )
+            expected = "{red}ERROR{nc}: message - description".format(**colors)
 
             assert expected == formatter.format(caplog.records[0])
 
@@ -95,7 +93,7 @@ def test_exception_under_verbose(self, caplog):
                 "{red}ERROR{nc}: description\n"
                 "{red}{line}{nc}\n"
                 "{stack_trace}"
-                "{red}{line}{nc}\n".format(
+                "{red}{line}{nc}".format(
                     line="-" * 60, stack_trace=stack_trace, **colors
                 )
             )
@@ -114,7 +112,7 @@ def test_tb_only(self, caplog):
                 "{red}ERROR{nc}: something\n"
                 "{red}{line}{nc}\n"
                 "{stack_trace}"
-                "{red}{line}{nc}\n".format(
+                "{red}{line}{nc}".format(
                     line="-" * 60, stack_trace=stack_trace, **colors
                 )
             )
@@ -136,7 +134,7 @@ def test_nested_exceptions(self, caplog):
                 "{red}ERROR{nc}: message - second: first\n"
                 "{red}{line}{nc}\n"
                 "{stack_trace}"
-                "{red}{line}{nc}\n".format(
+                "{red}{line}{nc}".format(
                     line="-" * 60, stack_trace=stack_trace, **colors
                 )
             )