|
from __future__ import annotations |
|
|
|
import http |
|
import logging |
|
import sys |
|
from copy import copy |
|
from typing import Literal |
|
|
|
import click |
|
|
|
TRACE_LOG_LEVEL = 5 |
|
|
|
|
|
class ColourizedFormatter(logging.Formatter): |
|
""" |
|
A custom log formatter class that: |
|
|
|
* Outputs the LOG_LEVEL with an appropriate color. |
|
* If a log call includes an `extra={"color_message": ...}` it will be used |
|
for formatting the output, instead of the plain text message. |
|
""" |
|
|
|
level_name_colors = { |
|
TRACE_LOG_LEVEL: lambda level_name: click.style(str(level_name), fg="blue"), |
|
logging.DEBUG: lambda level_name: click.style(str(level_name), fg="cyan"), |
|
logging.INFO: lambda level_name: click.style(str(level_name), fg="green"), |
|
logging.WARNING: lambda level_name: click.style(str(level_name), fg="yellow"), |
|
logging.ERROR: lambda level_name: click.style(str(level_name), fg="red"), |
|
logging.CRITICAL: lambda level_name: click.style(str(level_name), fg="bright_red"), |
|
} |
|
|
|
def __init__( |
|
self, |
|
fmt: str | None = None, |
|
datefmt: str | None = None, |
|
style: Literal["%", "{", "$"] = "%", |
|
use_colors: bool | None = None, |
|
): |
|
if use_colors in (True, False): |
|
self.use_colors = use_colors |
|
else: |
|
self.use_colors = sys.stdout.isatty() |
|
super().__init__(fmt=fmt, datefmt=datefmt, style=style) |
|
|
|
def color_level_name(self, level_name: str, level_no: int) -> str: |
|
def default(level_name: str) -> str: |
|
return str(level_name) |
|
|
|
func = self.level_name_colors.get(level_no, default) |
|
return func(level_name) |
|
|
|
def should_use_colors(self) -> bool: |
|
return True |
|
|
|
def formatMessage(self, record: logging.LogRecord) -> str: |
|
recordcopy = copy(record) |
|
levelname = recordcopy.levelname |
|
seperator = " " * (8 - len(recordcopy.levelname)) |
|
if self.use_colors: |
|
levelname = self.color_level_name(levelname, recordcopy.levelno) |
|
if "color_message" in recordcopy.__dict__: |
|
recordcopy.msg = recordcopy.__dict__["color_message"] |
|
recordcopy.__dict__["message"] = recordcopy.getMessage() |
|
recordcopy.__dict__["levelprefix"] = levelname + ":" + seperator |
|
return super().formatMessage(recordcopy) |
|
|
|
|
|
class DefaultFormatter(ColourizedFormatter): |
|
def should_use_colors(self) -> bool: |
|
return sys.stderr.isatty() |
|
|
|
|
|
class AccessFormatter(ColourizedFormatter): |
|
status_code_colours = { |
|
1: lambda code: click.style(str(code), fg="bright_white"), |
|
2: lambda code: click.style(str(code), fg="green"), |
|
3: lambda code: click.style(str(code), fg="yellow"), |
|
4: lambda code: click.style(str(code), fg="red"), |
|
5: lambda code: click.style(str(code), fg="bright_red"), |
|
} |
|
|
|
def get_status_code(self, status_code: int) -> str: |
|
try: |
|
status_phrase = http.HTTPStatus(status_code).phrase |
|
except ValueError: |
|
status_phrase = "" |
|
status_and_phrase = f"{status_code} {status_phrase}" |
|
if self.use_colors: |
|
|
|
def default(code: int) -> str: |
|
return status_and_phrase |
|
|
|
func = self.status_code_colours.get(status_code // 100, default) |
|
return func(status_and_phrase) |
|
return status_and_phrase |
|
|
|
def formatMessage(self, record: logging.LogRecord) -> str: |
|
recordcopy = copy(record) |
|
( |
|
client_addr, |
|
method, |
|
full_path, |
|
http_version, |
|
status_code, |
|
) = recordcopy.args |
|
status_code = self.get_status_code(int(status_code)) |
|
request_line = f"{method} {full_path} HTTP/{http_version}" |
|
if self.use_colors: |
|
request_line = click.style(request_line, bold=True) |
|
recordcopy.__dict__.update( |
|
{ |
|
"client_addr": client_addr, |
|
"request_line": request_line, |
|
"status_code": status_code, |
|
} |
|
) |
|
return super().formatMessage(recordcopy) |
|
|