From 712e62b9b07f690abbb40e089a17f4ddec6ba952 Mon Sep 17 00:00:00 2001
From: Aravind Vasudevan <aravindvasudev@google.com>
Date: Wed, 6 Sep 2023 17:25:58 +0000
Subject: [PATCH] logging: Use log.formatter for coloring logs

Bug: b/292704435
Change-Id: Iebdf8fb7666592dc5df2b36aae3185d1fc71bd66
Reviewed-on: https://gerrit-review.googlesource.com/c/git-repo/+/385514
Tested-by: Aravind Vasudevan <aravindvasudev@google.com>
Commit-Queue: Aravind Vasudevan <aravindvasudev@google.com>
Reviewed-by: Mike Frysinger <vapier@google.com>
---
 repo_logging.py            | 64 +++++++++++++++++++++-----------------
 tests/test_repo_logging.py | 36 ---------------------
 2 files changed, 36 insertions(+), 64 deletions(-)

diff --git a/repo_logging.py b/repo_logging.py
index e94af7df..7d050555 100644
--- a/repo_logging.py
+++ b/repo_logging.py
@@ -15,7 +15,7 @@
 """Logic for printing user-friendly logs in repo."""
 
 import logging
-from typing import Any, List
+from typing import List
 
 from color import Coloring
 
@@ -23,16 +23,7 @@ from color import Coloring
 SEPARATOR = "=" * 80
 
 
-class LogColoring(Coloring):
-    """Coloring outstream for logging."""
-
-    def __init__(self, config):
-        super().__init__(config, "logs")
-        self.error = self.colorer("error", fg="red")
-        self.warning = self.colorer("warn", fg="yellow")
-
-
-class ConfigMock:
+class _ConfigMock:
     """Default coloring config to use when Logging.config is not set."""
 
     def __init__(self):
@@ -42,28 +33,45 @@ class ConfigMock:
         return self.default_values.get(x, None)
 
 
+class _LogColoring(Coloring):
+    """Coloring outstream for logging."""
+
+    def __init__(self, config):
+        super().__init__(config, "logs")
+        self.error = self.colorer("error", fg="red")
+        self.warning = self.colorer("warn", fg="yellow")
+        self.levelMap = {
+            "WARNING": self.warning,
+            "ERROR": self.error,
+        }
+
+
+class _LogColoringFormatter(logging.Formatter):
+    """Coloring formatter for logging."""
+
+    def __init__(self, config=None, *args, **kwargs):
+        self.config = config if config else _ConfigMock()
+        self.colorer = _LogColoring(self.config)
+        super().__init__(*args, **kwargs)
+
+    def format(self, record):
+        """Formats |record| with color."""
+        msg = super().format(record)
+        colorer = self.colorer.levelMap.get(record.levelname)
+        return msg if not colorer else colorer(msg)
+
+
 class RepoLogger(logging.Logger):
     """Repo Logging Module."""
 
     def __init__(self, name: str, config=None, **kwargs):
         super().__init__(name, **kwargs)
-        self.config = config if config else ConfigMock()
-        self.colorer = LogColoring(self.config)
-
-    def error(self, msg: Any, *args, **kwargs):
-        """Print and aggregate error-level logs."""
-        colored_error = self.colorer.error(str(msg), *args)
-        super().error(colored_error, **kwargs)
-
-    def warning(self, msg: Any, *args, **kwargs):
-        """Print warning-level logs with coloring."""
-        colored_warning = self.colorer.warning(str(msg), *args)
-        super().warning(colored_warning, **kwargs)
+        handler = logging.StreamHandler()
+        handler.setFormatter(_LogColoringFormatter(config))
+        self.addHandler(handler)
 
     def log_aggregated_errors(self, errors: List[Exception]):
         """Print all aggregated logs."""
-        super().error(self.colorer.error(SEPARATOR))
-        super().error(
-            self.colorer.error("Repo command failed due to following errors:")
-        )
-        super().error("\n".join(map(str, errors)))
+        super().error(SEPARATOR)
+        super().error("Repo command failed due to following errors:")
+        super().error("\n".join(str(e) for e in errors))
diff --git a/tests/test_repo_logging.py b/tests/test_repo_logging.py
index b51e6270..52f251a7 100644
--- a/tests/test_repo_logging.py
+++ b/tests/test_repo_logging.py
@@ -20,42 +20,6 @@ from repo_logging import RepoLogger
 
 
 class TestRepoLogger(unittest.TestCase):
-    def test_error_logs_error(self):
-        """Test if error fn outputs logs."""
-        logger = RepoLogger(__name__)
-        result = None
-
-        def mock_handler(log):
-            nonlocal result
-            result = log.getMessage()
-
-        mock_out = mock.MagicMock()
-        mock_out.level = 0
-        mock_out.handle = mock_handler
-        logger.addHandler(mock_out)
-
-        logger.error("We're no strangers to love")
-
-        self.assertEqual(result, "We're no strangers to love")
-
-    def test_warning_logs_error(self):
-        """Test if warning fn outputs logs."""
-        logger = RepoLogger(__name__)
-        result = None
-
-        def mock_handler(log):
-            nonlocal result
-            result = log.getMessage()
-
-        mock_out = mock.MagicMock()
-        mock_out.level = 0
-        mock_out.handle = mock_handler
-        logger.addHandler(mock_out)
-
-        logger.warning("You know the rules and so do I (do I)")
-
-        self.assertEqual(result, "You know the rules and so do I (do I)")
-
     def test_log_aggregated_errors_logs_aggregated_errors(self):
         """Test if log_aggregated_errors outputs aggregated errors."""
         logger = RepoLogger(__name__)