Skip to content

Commit

Permalink
ansi_handler: support iterables in record message
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
Javagedes committed Aug 3, 2023
1 parent 85b899a commit ae1663c
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 7 deletions.
9 changes: 5 additions & 4 deletions edk2toollib/log/ansi_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -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


Expand Down Expand Up @@ -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)

Check warning on line 444 in edk2toollib/log/ansi_handler.py

View check run for this annotation

Codecov / codecov/patch

edk2toollib/log/ansi_handler.py#L443-L444

Added lines #L443 - L444 were not covered by tests
self.handleError(record)
34 changes: 31 additions & 3 deletions tests.unit/test_ansi_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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")
Expand Down Expand Up @@ -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

0 comments on commit ae1663c

Please sign in to comment.