From ae1663ccdc9624e84ece2e1945eeb9ef4fa0c9b6 Mon Sep 17 00:00:00 2001
From: Joey Vagedes <joey.vagedes@gmail.com>
Date: Thu, 3 Aug 2023 09:43:20 -0700
Subject: [PATCH] ansi_handler: support iterables in record message

Previously, the ansi_handler only accepted strings in the logging record
message due to a bug in how the message was retrieved from the record.
This change fixes this, allowing any object that `str()` works on to be
used as a message in the logging record.
---
 edk2toollib/log/ansi_handler.py |  9 +++++----
 tests.unit/test_ansi_handler.py | 34 ++++++++++++++++++++++++++++++---
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/edk2toollib/log/ansi_handler.py b/edk2toollib/log/ansi_handler.py
index 76f88fc7..1224dee6 100644
--- a/edk2toollib/log/ansi_handler.py
+++ b/edk2toollib/log/ansi_handler.py
@@ -203,7 +203,7 @@ def __init__(self, msg="", use_azure=False):
     def format(self, record):
         """Formats the given record and returns it."""
         levelname = record.levelname
-        org_message = record.msg
+        message = record.getMessage()
 
         if not self.use_azure and levelname in ColoredFormatter.COLORS:
             # just color the level name
@@ -212,7 +212,7 @@ def format(self, record):
             # otherwise color the wholes message
             else:
                 levelname_color = get_ansi_string(ColoredFormatter.COLORS[levelname]) + levelname
-                record.msg += get_ansi_string()
+                message += get_ansi_string()
             record.levelname = levelname_color
 
         if self.use_azure and levelname in ColoredFormatter.AZURE_COLORS:
@@ -223,7 +223,7 @@ def format(self, record):
         result = logging.Formatter.format(self, record)
 
         record.levelname = levelname
-        record.msg = org_message
+        record.msg = message
         return result
 
 
@@ -440,5 +440,6 @@ def emit(self, record):
             self.write(str(msg))
             self.write(self.terminator)
             self.flush()
-        except Exception:
+        except Exception as e:
+            print(e)
             self.handleError(record)
diff --git a/tests.unit/test_ansi_handler.py b/tests.unit/test_ansi_handler.py
index 346a38ac..4298088f 100644
--- a/tests.unit/test_ansi_handler.py
+++ b/tests.unit/test_ansi_handler.py
@@ -5,10 +5,10 @@
 #
 # SPDX-License-Identifier: BSD-2-Clause-Patent
 ##
-import unittest
 import logging
-from edk2toollib.log.ansi_handler import ColoredFormatter
-from edk2toollib.log.ansi_handler import ColoredStreamHandler
+import unittest
+
+from edk2toollib.log.ansi_handler import ColoredFormatter, ColoredStreamHandler
 
 try:
     from StringIO import StringIO
@@ -26,6 +26,15 @@ class AnsiHandlerTest(unittest.TestCase):
     record2 = logging.makeLogRecord({"name": "", "level": logging.INFO, "levelno": logging.INFO,
                                      "levelname": "INFO", "path": "test_path", "lineno": 0,
                                      "msg": "Test message"})
+    record3 = logging.makeLogRecord({"name": "", "level": logging.ERROR, "levelno": logging.ERROR,
+                                     "levelname": "ERROR", "path": "test_path", "lineno": 0,
+                                     "msg": ['Logging', 'A', 'List']})
+    record4 = logging.makeLogRecord({"name": "", "level": logging.ERROR, "levelno": logging.ERROR,
+                                     "levelname": "ERROR", "path": "test_path", "lineno": 0,
+                                     "msg": ('Logging', 'A', 'Tuple')})
+    record5 = logging.makeLogRecord({"name": "", "level": logging.ERROR, "levelno": logging.ERROR,
+                                     "levelname": "ERROR", "path": "test_path", "lineno": 0,
+                                     "msg": iter(('Logging', 'A', 'Iter'))})
 
     def test_colored_formatter_init(self):
         formatter = ColoredFormatter("%(levelname)s - %(message)s")
@@ -82,3 +91,22 @@ def test_color_handler_not_strip_ansi(self):
             if CSI in line:
                 found_csi = True
         self.assertTrue(found_csi, "We are supposed to to have found an ANSI control character %s" % lines)
+
+    def test_ansi_handler_with_list(self):
+        """Tests that the ANSI handler can handle Iterables in the message."""
+        stream = StringIO()
+        formatter = ColoredFormatter("%(levelname)s - %(message)s")
+        handler = ColoredStreamHandler(stream, strip=False, convert=False)
+        handler.setFormatter(formatter)
+        handler.setLevel(logging.INFO)
+
+        handler.emit(AnsiHandlerTest.record3)
+        handler.emit(AnsiHandlerTest.record4)
+        handler.emit(AnsiHandlerTest.record5)
+        handler.flush()
+
+        stream.seek(0)
+        lines = stream.readlines()
+        CSI = '\033[31m'
+        for line in lines:
+            assert CSI in line