Ran update-copyright.py.

[pygrader.git] / pygrader / color.py

diff --git a/pygrader/color.py b/pygrader/color.py
index 741fd34b363aa8618d1455595a7ade0b44cb0f22..5c4d44a006d861588f31b75ccd633a68f585b0c3 100644 (file)
--- a/pygrader/color.py
+++ b/pygrader/color.py
@@ -1,5 +1,20 @@
-# Copyright
-
+# Copyright (C) 2012 W. Trevor King <wking@tremily.us>
+#
+# This file is part of pygrader.
+#
+# pygrader is free software: you can redistribute it and/or modify it under the
+# terms of the GNU General Public License as published by the Free Software
+# Foundation, either version 3 of the License, or (at your option) any later
+# version.
+#
+# pygrader is distributed in the hope that it will be useful, but WITHOUT ANY
+# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
+# A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along with
+# pygrader.  If not, see <http://www.gnu.org/licenses/>.
+
+import logging as _logging
 import sys as _sys
 
 
@@ -8,6 +23,8 @@ _COLORS = [
     'black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white']
 
 USE_COLOR = True
+GOOD_DEBUG = _logging.DEBUG + 2
+BAD_DEBUG = _logging.DEBUG + 4
 
 
 def standard_colors(use_color=None):
@@ -83,3 +100,26 @@ def write_color(string, color=None, stream=None):
         stream = _sys.stdout
     stream.write(color_string(string=string, color=color))
     stream.flush()
+
+
+class ColoredFormatter (_logging.Formatter):
+    def __init__(self, *args, **kwargs):
+        super(ColoredFormatter, self).__init__(*args, **kwargs)
+        self.colored = None  # `None` to use USE_COLOR; True/False to override
+
+    def format(self, record):
+        s = super(ColoredFormatter, self).format(record)
+        if self.colored or (self.colored is None and USE_COLOR):
+            highlight,lowlight,good,bad = standard_colors()
+            if record.levelno <= _logging.DEBUG:
+                color = lowlight
+            elif record.levelno <= GOOD_DEBUG:
+                color = good
+            elif record.levelno <= BAD_DEBUG:
+                color = bad
+            elif record.levelno <= _logging.INFO:
+                color = highlight
+            else:
+                color = bad
+            return color_string(string=s, color=color)
+        return s