From 5c71e321b4d168444721cced85e7a86c9bacb5ad Mon Sep 17 00:00:00 2001
From: Elena ``of Valhalla'' Grandi <valhalla@trueelena.org>
Date: Tue, 2 Aug 2022 09:52:40 +0200
Subject: Add optional support for coloredlogs

---
 CHANGELOG.rst         |  3 +++
 hazwaz/command.py     | 31 +++++++++++++++++++++----------
 tests/test_command.py | 44 +++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 65 insertions(+), 13 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index b65d52f..7d00dac 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -7,6 +7,9 @@ Unreleased
 
 * Add helper code for testing. See
   https://hazwaz.trueelena.org/testing.html
+* Added optional support for `coloredlogs
+  <https://coloredlogs.readthedocs.io/>`_ (if available it will be used
+  by default unless explicitely disabled).
 
 0.0.1
 =====
diff --git a/hazwaz/command.py b/hazwaz/command.py
index df8c186..c2dbe95 100644
--- a/hazwaz/command.py
+++ b/hazwaz/command.py
@@ -2,7 +2,10 @@ import argparse
 import logging
 import typing
 
-logger = logging.getLogger(__name__)
+try:
+    import coloredlogs  # type: ignore
+except ModuleNotFoundError:
+    coloredlogs = None
 
 
 def _get_first_docstring_line(obj):
@@ -43,6 +46,10 @@ class MainCommand:
     """
     The format passed to logging.Formatter.
     """
+    coloredlogs: bool = True
+    """
+    Whether coloredlogs is used (if available)
+    """
 
     def __init__(self):
         desc = _get_first_docstring_line(self)
@@ -95,18 +102,22 @@ class MainCommand:
         )
 
     def setup_logging(self):
-        logger = logging.getLogger()
-        handler = logging.StreamHandler()
-        formatter = logging.Formatter(self.logformat)
-        handler.setFormatter(formatter)
-        logger.addHandler(handler)
-
         if getattr(self.args, "debug", False):
-            logger.setLevel(logging.DEBUG)
+            level = logging.DEBUG
         elif getattr(self.args, "verbose", False):
-            logger.setLevel(logging.INFO)
+            level = logging.INFO
+        else:
+            level = logging.WARNING
+
+        if self.coloredlogs and coloredlogs:
+            coloredlogs.install(level=level, fmt=self.logformat)
         else:
-            logger.setLevel(logging.WARNING)
+            logger = logging.getLogger()
+            handler = logging.StreamHandler()
+            formatter = logging.Formatter(self.logformat)
+            handler.setFormatter(formatter)
+            logger.addHandler(handler)
+            logger.setLevel(level)
 
     def run(self):
         """
diff --git a/tests/test_command.py b/tests/test_command.py
index 1d1402c..667fef6 100644
--- a/tests/test_command.py
+++ b/tests/test_command.py
@@ -155,8 +155,9 @@ class testCommand(hazwaz.unittest.HazwazTestCase):
         stream = self.run_with_argv(cmd, ["mycommand"])
         self.assertEqual(stream["stdout"].getvalue(), cmd_help)
 
-    def test_logging_regular(self):
+    def test_logging_regular_coloredlogs(self):
         cmd = MyCommand()
+        cmd.coloredlogs = True
         with self.assertLogs():
             stream = self.run_with_argv(cmd, [
                 "mycommand",
@@ -165,8 +166,9 @@ class testCommand(hazwaz.unittest.HazwazTestCase):
         log_lines = stream["stderr"].getvalue().strip().split("\n")
         self.assertEqual(len(log_lines), 1)
 
-    def test_logging_verbose(self):
+    def test_logging_verbose_coloredlogs(self):
         cmd = MyCommand()
+        cmd.coloredlogs = True
         with self.assertLogs():
             stream = self.run_with_argv(cmd, [
                 "mycommand",
@@ -176,8 +178,44 @@ class testCommand(hazwaz.unittest.HazwazTestCase):
         log_lines = stream["stderr"].getvalue().strip().split("\n")
         self.assertEqual(len(log_lines), 2)
 
-    def test_logging_debug(self):
+    def test_logging_debug_coloredlogs(self):
         cmd = MyCommand()
+        cmd.coloredlogs = True
+        with self.assertLogs():
+            stream = self.run_with_argv(cmd, [
+                "mycommand",
+                "--debug",
+                "loggingsubcommand",
+            ])
+        log_lines = stream["stderr"].getvalue().strip().split("\n")
+        self.assertEqual(len(log_lines), 3)
+
+    def test_logging_regular_no_coloredlogs(self):
+        cmd = MyCommand()
+        cmd.coloredlogs = False
+        with self.assertLogs():
+            stream = self.run_with_argv(cmd, [
+                "mycommand",
+                "loggingsubcommand",
+            ])
+        log_lines = stream["stderr"].getvalue().strip().split("\n")
+        self.assertEqual(len(log_lines), 1)
+
+    def test_logging_verbose_no_coloredlogs(self):
+        cmd = MyCommand()
+        cmd.coloredlogs = False
+        with self.assertLogs():
+            stream = self.run_with_argv(cmd, [
+                "mycommand",
+                "--verbose",
+                "loggingsubcommand",
+            ])
+        log_lines = stream["stderr"].getvalue().strip().split("\n")
+        self.assertEqual(len(log_lines), 2)
+
+    def test_logging_debug_no_coloredlogs(self):
+        cmd = MyCommand()
+        cmd.coloredlogs = False
         with self.assertLogs():
             stream = self.run_with_argv(cmd, [
                 "mycommand",
-- 
cgit v1.2.3