Add utils.setup_logging to help set up logging outside of Client.run

discord/utils.py

@@ -63,9 +63,11 @@ from operator import attrgetter
 from urllib.parse import urlencode
 import json
 import re
+import os
 import sys
 import types
 import warnings
+import logging
 
 import yarl
 
@@ -91,6 +93,7 @@ __all__ = (
     'as_chunks',
     'format_dt',
     'MISSING',
+    'setup_logging',
 )
 
 DISCORD_EPOCH = 1420070400000
@@ -1180,3 +1183,117 @@ def format_dt(dt: datetime.datetime, /, style: Optional[TimestampStyle] = None)
     if style is None:
         return f'<t:{int(dt.timestamp())}>'
     return f'<t:{int(dt.timestamp())}:{style}>'
+
+
+def stream_supports_colour(stream: Any) -> bool:
+    is_a_tty = hasattr(stream, 'isatty') and stream.isatty()
+    if sys.platform != 'win32':
+        return is_a_tty
+
+    # ANSICON checks for things like ConEmu
+    # WT_SESSION checks if this is Windows Terminal
+    # VSCode built-in terminal supports colour too
+    return is_a_tty and ('ANSICON' in os.environ or 'WT_SESSION' in os.environ or os.environ.get('TERM_PROGRAM') == 'vscode')
+
+
+class _ColourFormatter(logging.Formatter):
+
+    # ANSI codes are a bit weird to decipher if you're unfamiliar with them, so here's a refresher
+    # It starts off with a format like \x1b[XXXm where XXX is a semicolon separated list of commands
+    # The important ones here relate to colour.
+    # 30-37 are black, red, green, yellow, blue, magenta, cyan and white in that order
+    # 40-47 are the same except for the background
+    # 90-97 are the same but "bright" foreground
+    # 100-107 are the same as the bright ones but for the background.
+    # 1 means bold, 2 means dim, 0 means reset, and 4 means underline.
+
+    LEVEL_COLOURS = [
+        (logging.DEBUG, '\x1b[40;1m'),
+        (logging.INFO, '\x1b[34;1m'),
+        (logging.WARNING, '\x1b[33;1m'),
+        (logging.ERROR, '\x1b[31m'),
+        (logging.CRITICAL, '\x1b[41m'),
+    ]
+
+    FORMATS = {
+        level: logging.Formatter(
+            f'\x1b[30;1m%(asctime)s\x1b[0m {colour}%(levelname)-8s\x1b[0m \x1b[35m%(name)s\x1b[0m %(message)s',
+            '%Y-%m-%d %H:%M:%S',
+        )
+        for level, colour in LEVEL_COLOURS
+    }
+
+    def format(self, record):
+        formatter = self.FORMATS.get(record.levelno)
+        if formatter is None:
+            formatter = self.FORMATS[logging.DEBUG]
+
+        # Override the traceback to always print in red
+        if record.exc_info:
+            text = formatter.formatException(record.exc_info)
+            record.exc_text = f'\x1b[31m{text}\x1b[0m'
+
+        output = formatter.format(record)
+
+        # Remove the cache layer
+        record.exc_text = None
+        return output
+
+
+def setup_logging(
+    *,
+    handler: logging.Handler = MISSING,
+    formatter: logging.Formatter = MISSING,
+    level: int = MISSING,
+    root: bool = True,
+) -> None:
+    """A helper function to setup logging.
+
+    This is superficially similar to :func:`logging.basicConfig` but
+    uses different defaults and a colour formatter if the stream can
+    display colour.
+
+    This is used by the :class:`~discord.Client` to set up logging
+    if ``log_handler`` is not ``None``.
+
+    .. versionadded:: 2.0
+
+    Parameters
+    -----------
+    handler: :class:`logging.Handler`
+        The log handler to use for the library's logger.
+
+        The default log handler if not provided is :class:`logging.StreamHandler`.
+    formatter: :class:`logging.Formatter`
+        The formatter to use with the given log handler. If not provided then it
+        defaults to a colour based logging formatter (if available). If colour
+        is not available then a simple logging formatter is provided.
+    level: :class:`int`
+        The default log level for the library's logger. Defaults to ``logging.INFO``.
+    root: :class:`bool`
+        Whether to set up the root logger rather than the library logger.
+        Unlike the default for :class:`~discord.Client`, this defaults to ``True``.
+    """
+
+    if level is MISSING:
+        level = logging.INFO
+
+    if handler is MISSING:
+        handler = logging.StreamHandler()
+
+    if formatter is MISSING:
+        if isinstance(handler, logging.StreamHandler) and stream_supports_colour(handler.stream):
+            formatter = _ColourFormatter()
+        else:
+            dt_fmt = '%Y-%m-%d %H:%M:%S'
+            formatter = logging.Formatter('[{asctime}] [{levelname:<8}] {name}: {message}', dt_fmt, style='{')
+
+    if root:
+        logger = logging.getLogger()
+    else:
+        library, _, _ = __name__.partition('.')
+        logger = logging.getLogger(library)
+
+    handler.setFormatter(formatter)
+    logger.setLevel(level)
+    logger.addHandler(handler)