add docstrings and make it possible to just set a subset of the color "palette" (dict updated).

1 files changed, 19 insertions, 10 deletions

diff --git a/logging_ext.py b/logging_ext.py
index a3624ae..37a2168 100644
--- a/logging_ext.py
+++ b/logging_ext.py
@@ -35,7 +35,7 @@ class ColorFormatter(logging.Formatter):
     By default, colorize CRITICAL and ERROR in red, WARNING in orange
     and INFO in yellow.
 
-    self.colors is customizable via the constructor.
+    self.colors is customizable via the 'color' constructor argument (dictionnary).
 
     self.colorfilters is a list of functions that get the LogRecord
     and return a color name or None.
@@ -44,13 +44,15 @@ class ColorFormatter(logging.Formatter):
     def __init__(self, fmt=None, datefmt=None, colors=None):
         logging.Formatter.__init__(self, fmt, datefmt)
         self.colorfilters = []
-        self.colors = colors or {'CRITICAL': 'red',
-                                 'ERROR': 'red',
-                                 'WARNING': 'magenta',
-                                 'INFO': 'yellow',
-                                 }
-        assert isinstance(self.colors, dict)
-        
+        self.colors = {'CRITICAL': 'red',
+                       'ERROR': 'red',
+                       'WARNING': 'magenta',
+                       'INFO': 'yellow',
+                       }
+        if colors is not None:
+            assert isinstance(colors, dict)
+            self.colors.update(colors)            
+                               
     def format(self, record):
         msg = logging.Formatter.format(self, record)
         if record.levelname in self.colors:
@@ -63,12 +65,19 @@ class ColorFormatter(logging.Formatter):
                     return colorize_ansi(msg, color)
         return msg
 
-def set_color_formatter(logger=None):
+def set_color_formatter(logger=None, **kw):
+    """
+    Install a color formatter on the 'logger'. If not given, it will
+    defaults to the default logger.
+
+    Any additional keyword will be passed as-is to the ColorFormatter
+    constructor.
+    """
     if logger is None:
         logger = logging.getLogger()
         if not logger.handlers:
             logging.basicConfig()
     format_msg = logger.handlers[0].formatter._fmt
-    fmt = ColorFormatter(format_msg)
+    fmt = ColorFormatter(format_msg, **kw)
     fmt.colorfilters.append(xxx_cyan)
     logger.handlers[0].setFormatter(fmt)