Make colors customizable

microsoft · Mar 22, 2021 · b3f3c30 · b3f3c30
1 parent fd36947
commit b3f3c30
Show file tree

Hide file tree

Showing 7 changed files with 62 additions and 42 deletions.
diff --git a/examples/exapp2 b/examples/exapp2
@@ -125,6 +125,24 @@ def sample_json_handler():
     return result
 
 
+def sample_logger_handler():
+    """
+    Print logs to stderr.
+    """
+    print("""This is a demo for logging. The logging level can be controlled with:
+  --only-show-errors: Show ERROR logs and above
+  <default>:          Show WARNING logs and above
+  --verbose:          Show INFO logs and above
+  --debug:            Show DEBUG logs and above""")
+    from knack.log import get_logger
+    logger = get_logger(__name__)
+    logger.debug("This is a debug log entry.")
+    logger.info("This is a info log entry.")
+    logger.warning("This is a warning log entry.")
+    logger.error("This is a error log entry.")
+    logger.critical("This is a critical log entry.")
+
+
 def hello_command_handler(greetings=None):
     """
     Say "Hello World!" and my warm greetings
@@ -160,6 +178,7 @@ class MyCommandsLoader(CLICommandsLoader):
         with CommandGroup(self, '', '__main__#{}') as g:
             g.command('hello', 'hello_command_handler', confirmation=True)
             g.command('sample-json', 'sample_json_handler')
+            g.command('sample-logger', 'sample_logger_handler')
         with CommandGroup(self, 'abc', '__main__#{}') as g:
             g.command('list', 'abc_list_command_handler')
             g.command('show', 'abc_show_command_handler')

diff --git a/knack/deprecation.py b/knack/deprecation.py
@@ -3,9 +3,10 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-from .util import StatusTag
+from .util import StatusTag, color_map
 
 DEFAULT_DEPRECATED_TAG = '[Deprecated]'
+_config_key = 'deprecation'
 
 
 def resolve_deprecate_info(cli_ctx, name):
@@ -88,7 +89,7 @@ def _default_get_message(self):
             cli_ctx=cli_ctx,
             object_type=object_type,
             target=target,
-            color='yellow',
+            color=color_map[_config_key],
             tag_func=tag_func or (lambda _: DEFAULT_DEPRECATED_TAG),
             message_func=message_func or _default_get_message
         )

diff --git a/knack/experimental.py b/knack/experimental.py
@@ -3,10 +3,11 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-from .util import StatusTag, status_tag_messages
+from .util import StatusTag, status_tag_messages, color_map
 
 _EXPERIMENTAL_TAG = '[Experimental]'
 _experimental_kwarg = 'experimental_info'
+_config_key = 'experimental'
 
 
 def resolve_experimental_info(cli_ctx, name):
@@ -50,13 +51,13 @@ def __init__(self, cli_ctx, object_type='', target=None, tag_func=None, message_
         """
 
         def _default_get_message(self):
-            return status_tag_messages['experimental'].format("This " + self.object_type)
+            return status_tag_messages[_config_key].format("This " + self.object_type)
 
         super().__init__(
             cli_ctx=cli_ctx,
             object_type=object_type,
             target=target,
-            color='cyan',
+            color=color_map[_config_key],
             tag_func=tag_func or (lambda _: _EXPERIMENTAL_TAG),
             message_func=message_func or _default_get_message
         )
@@ -67,7 +68,7 @@ class ImplicitExperimentalItem(ExperimentalItem):
     def __init__(self, **kwargs):
 
         def get_implicit_experimental_message(self):
-            return status_tag_messages['experimental'].format("Command group '{}'".format(self.target))
+            return status_tag_messages[_config_key].format("Command group '{}'".format(self.target))
 
         kwargs.update({
             'tag_func': lambda _: '',

diff --git a/knack/log.py b/knack/log.py
@@ -7,7 +7,7 @@
 import logging
 from enum import IntEnum
 
-from .util import CtxTypeError, ensure_dir, CLIError
+from .util import CtxTypeError, ensure_dir, CLIError, color_map
 from .events import EVENT_PARSER_GLOBAL_CREATE
 
 
@@ -44,27 +44,11 @@ def get_logger(module_name=None):
 
 
 class _CustomStreamHandler(logging.StreamHandler):
-    COLOR_MAP = None
 
     @classmethod
-    def get_color_wrapper(cls, level):
-        if not cls.COLOR_MAP:
-            import colorama
-
-            def _color_wrapper(color_marker):
-                def wrap_msg_with_color(msg):
-                    return '{}{}{}'.format(color_marker, msg, colorama.Style.RESET_ALL)
-                return wrap_msg_with_color
-
-            cls.COLOR_MAP = {
-                logging.CRITICAL: _color_wrapper(colorama.Fore.LIGHTRED_EX),
-                logging.ERROR: _color_wrapper(colorama.Fore.LIGHTRED_EX),
-                logging.WARNING: _color_wrapper(colorama.Fore.YELLOW),
-                logging.INFO: _color_wrapper(colorama.Fore.GREEN),
-                logging.DEBUG: _color_wrapper(colorama.Fore.CYAN)
-            }
-
-        return cls.COLOR_MAP.get(level, None)
+    def wrap_with_color(cls, level_name, msg):
+        color_marker = color_map[level_name.lower()]
+        return '{}{}{}'.format(color_marker, msg, color_map['reset'])
 
     def __init__(self, log_level_config, log_format, enable_color):
         logging.StreamHandler.__init__(self)
@@ -76,7 +60,7 @@ def format(self, record):
         msg = logging.StreamHandler.format(self, record)
         if self.enable_color:
             try:
-                msg = self.get_color_wrapper(record.levelno)(msg)
+                msg = self.wrap_with_color(record.levelname, msg)
             except KeyError:
                 pass
         return msg

diff --git a/knack/preview.py b/knack/preview.py
@@ -3,10 +3,11 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-from .util import StatusTag, status_tag_messages
+from .util import StatusTag, status_tag_messages, color_map
 
 _PREVIEW_TAG = '[Preview]'
 _preview_kwarg = 'preview_info'
+_config_key = 'preview'
 
 
 def resolve_preview_info(cli_ctx, name):
@@ -50,13 +51,13 @@ def __init__(self, cli_ctx, object_type='', target=None, tag_func=None, message_
         """
 
         def _default_get_message(self):
-            return status_tag_messages['preview'].format("This " + self.object_type)
+            return status_tag_messages[_config_key].format("This " + self.object_type)
 
         super().__init__(
             cli_ctx=cli_ctx,
             object_type=object_type,
             target=target,
-            color='cyan',
+            color=color_map[_config_key],
             tag_func=tag_func or (lambda _: _PREVIEW_TAG),
             message_func=message_func or _default_get_message
         )
@@ -67,7 +68,7 @@ class ImplicitPreviewItem(PreviewItem):
     def __init__(self, **kwargs):
 
         def get_implicit_preview_message(self):
-            return status_tag_messages['preview'].format("Command group '{}'".format(self.target))
+            return status_tag_messages[_config_key].format("Command group '{}'".format(self.target))
 
         kwargs.update({
             'tag_func': lambda _: '',

diff --git a/knack/util.py b/knack/util.py
@@ -12,13 +12,26 @@
 NO_COLOR_VARIABLE_NAME = 'KNACK_NO_COLOR'
 
 # Override these values to customize the status message.
-# The message should contain a placeholder indicating the subject (like 'This command group', 'Commend group xxx').
+# The message should contain a placeholder indicating the subject (like 'This command group', 'Command group xxx').
 # (A dict is used to avoid the "from A import B" pitfall that creates a copy of the imported B.)
 status_tag_messages = {
     'preview': "{} is in preview. It may be changed/removed in a future release.",
     'experimental': "{} is experimental and under development."
 }
 
+# https://docs.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences
+color_map = {
+    'reset': '\x1b[0m',  # Default
+    'preview': '\x1b[36m',  # Foreground Cyan
+    'experimental': '\x1b[36m',  # Foreground Cyan
+    'deprecation': '\x1b[33m',   # Foreground Yellow
+    'critical': '\x1b[41m',  # Background Red
+    'error': '\x1b[91m',  # Bright Foreground Red
+    'warning': '\x1b[33m',  # Foreground Yellow
+    'info': '\x1b[32m',  # Foreground Green
+    'debug': '\x1b[36m',  # Foreground Cyan
+}
+
 
 class CommandResultItem(object):  # pylint: disable=too-few-public-methods
     def __init__(self, result, table_transformer=None, is_query_active=False,
@@ -48,18 +61,16 @@ def __init__(self, obj):
 class ColorizedString(object):
 
     def __init__(self, message, color):
-        import colorama
         self._message = message
-        self._color = getattr(colorama.Fore, color.upper(), None)
+        self._color = color
 
     def __len__(self):
         return len(self._message)
 
     def __str__(self):
-        import colorama
         if not self._color:
             return self._message
-        return self._color + self._message + colorama.Fore.RESET
+        return self._color + self._message + color_map['reset']
 
 
 class StatusTag(object):

diff --git a/tests/test_log.py b/tests/test_log.py
@@ -9,7 +9,6 @@
 except ImportError:
     from unittest import mock
 import logging
-import colorama
 
 from knack.events import EVENT_PARSER_GLOBAL_CREATE, EVENT_INVOKER_PRE_CMD_TBL_CREATE
 from knack.log import CLILogging, get_logger, CLI_LOGGER_NAME, _CustomStreamHandler
@@ -165,15 +164,19 @@ def test_get_console_log_formats(self):
 
 
 class TestCustomStreamHandler(unittest.TestCase):
-    expectation = {logging.CRITICAL: colorama.Fore.LIGHTRED_EX, logging.ERROR: colorama.Fore.LIGHTRED_EX,
-                   logging.WARNING: colorama.Fore.YELLOW, logging.INFO: colorama.Fore.GREEN,
-                   logging.DEBUG: colorama.Fore.CYAN}
+    expectation = {
+        'critical': '\x1b[41m',  # Background Red
+        'error': '\x1b[91m',  # Bright Foreground Red
+        'warning': '\x1b[33m',  # Foreground Yellow
+        'info': '\x1b[32m',  # Foreground Green
+        'debug': '\x1b[36m',  # Foreground Cyan
+    }
 
     def test_get_color_wrapper(self):
         for level, prefix in self.expectation.items():
-            message = _CustomStreamHandler.get_color_wrapper(level)('test')
+            message = _CustomStreamHandler.wrap_with_color(level, 'test')
             self.assertTrue(message.startswith(prefix))
-            self.assertTrue(message.endswith(colorama.Style.RESET_ALL))
+            self.assertTrue(message.endswith('\x1b[0m'))
 
 
 if __name__ == '__main__':