Make the colored logger best friends with celery :]

colored_logger.py

@@ -20,7 +20,8 @@ intro_words = ["magic", "MaGiC", "voodoo", "sorcery", "wizardry", "witchery", "f
 intro_adjectives = ["colourful", "evil", "fairy", "random"]
 intro_colors = [34, 36, 32, 33, 35]
 
-decorate_info = {
+
+logger_methods_colors = {
     # method names and their desired highlight colors
     "debug": "1;90",
     "info": "1;32",
@@ -32,7 +33,10 @@ decorate_info = {
 
 # Adjusting handlers: Python logging.StreamHandler targets stderr by default.
 # Django default logging configuration doesn't change it, too.
-# Existing loggers first.
+# New loggers (they will be initially created with the proper stream).
+original_constructor = logging.StreamHandler.__init__  # must grab a reference outside of the lambda
+logging.StreamHandler.__init__ = lambda zelf, stream=None: original_constructor(zelf, stream if stream != sys.stderr else sys.stdout)
+# Existing loggers
 patched_objects = set()
 for logger in logging.Logger.manager.loggerDict.values():
     for handler in getattr(logger, "handlers", list()):
@@ -40,45 +44,51 @@ for logger in logging.Logger.manager.loggerDict.values():
             if handler.stream == sys.stderr:
                 handler.stream = sys.stdout
                 patched_objects.add(handler)
-# Future loggers (they should be initially created with the proper stream).
-original_constructor = logging.StreamHandler.__init__  # must grab a reference outside of the lambda
-logging.StreamHandler.__init__ = lambda zelf, stream=None: original_constructor(zelf, stream or sys.stdout)
-# the isatty() stuff below solved some problems in the past, I'm not sure anymore which exactly
+
+
+# the isatty() stuff below solves the problem with IDE streams not being recognized as TTY
+# normally sys.stdout is sufficient, but some stubborn stuff like Celery explicitly checks sys.stderr even if it doesn't write there
+for stream in [sys.stdout, sys.stderr]:
     try:
-    stdout_tty = sys.stdout.isatty()
+        isatty = stream.isatty()
     except Exception:
-    stdout_tty = False
-if not stdout_tty:
-    patched_objects.add(sys.stdout.isatty)
-    sys.stdout.isatty = lambda: True  # stdout is now a TTY no matter what
+        isatty = False
+    if not isatty:
+        patched_objects.add(stream)
+        stream.isatty = lambda: True  # stdout is now a TTY no matter what
 
-def _apply_color_to_msg(msg, color, extras):
-    return "".join([color_set.format(color), extras, str(msg), color_reset])
+
+def colorize_string(msg, color="1;32", extras=""):
+    return str().join([color_set.format(color), extras, str(msg), color_reset])
 
 
 def decorate_for_tty(method, color="1;39", extras=""):
-    """ it also knows how to decorate bound methods, but it is not used anymore """
+    """
+    decorate a method, replacing its first argument (must be str!) with the same, but colorized
+    (it also knows how to decorate bound methods)
+    """
     method_is_bound = bool(getattr(method, "__self__", None))  # existing loggers cause bound methods
-    msg_index = 0 if method_is_bound else 1  # message argument is located there
+    msg_arg_index = 0 if method_is_bound else 1  # message argument is located there
 
     def wrapper(*args, **kwargs):
         args = list(args)  # it's an immutable tuple when received
-        args[msg_index] = _apply_color_to_msg(args[msg_index], color, extras)
+        args[msg_arg_index] = colorize_string(args[msg_arg_index], color, extras)
         return method(*args, **kwargs)
 
     return wrapper
 
+
+# decorate classes
 for entity in [logging.Logger]:
-    for method_name, color_str in decorate_info.items():
-        if isinstance(entity, logging.Logger) or entity is logging.Logger:
+    for method_name, color_str in logger_methods_colors.items():
+        if isinstance(entity, logging.PlaceHolder):
+            continue  # ignore the PlaceHolder object, it's not functional and not a public API
         new_method = decorate_for_tty(getattr(entity, method_name), color_str)
         setattr(entity, method_name, new_method)
         patched_objects.add(entity)
-        else:
-            pass  # that should be a placeholder object (non-public API), ignoring it
 
 
-# useless foo
+# fancy useless stuff
 intro_seq = zip(itertools.cycle(intro_colors), random.choice(intro_words))
 intro_msg = "{adj} {word}{reset}".format(
     adj=random.choice(intro_adjectives).capitalize(),
@@ -94,6 +104,8 @@ sys.stdout.write("{intro_msg}{c}: patched {n} object{s}{cr}\n".format(
     s="" if len(patched_objects) is 1 else "s",
     cr=color_reset,
 ))
-logging.getLogger("pydevutils").debug("Logging patches: {}.".format(", ".join(str(e) for e in patched_objects) or "none"))
+
+# debug with the following line:
+# sys.stdout.write("Patched objects: {}\n".format(", ".join(str(e) for e in patched_objects) or "none"))