logging.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119
  1. from __future__ import annotations
  2. import http
  3. import logging
  4. import sys
  5. from copy import copy
  6. from typing import Literal
  7. import click
  8. TRACE_LOG_LEVEL = 5
  9. class ColourizedFormatter(logging.Formatter):
  10. """
  11. A custom log formatter class that:
  12. * Outputs the LOG_LEVEL with an appropriate color.
  13. * If a log call includes an `extras={"color_message": ...}` it will be used
  14. for formatting the output, instead of the plain text message.
  15. """
  16. level_name_colors = {
  17. TRACE_LOG_LEVEL: lambda level_name: click.style(str(level_name), fg="blue"),
  18. logging.DEBUG: lambda level_name: click.style(str(level_name), fg="cyan"),
  19. logging.INFO: lambda level_name: click.style(str(level_name), fg="green"),
  20. logging.WARNING: lambda level_name: click.style(str(level_name), fg="yellow"),
  21. logging.ERROR: lambda level_name: click.style(str(level_name), fg="red"),
  22. logging.CRITICAL: lambda level_name: click.style(
  23. str(level_name), fg="bright_red"
  24. ),
  25. }
  26. def __init__(
  27. self,
  28. fmt: str | None = None,
  29. datefmt: str | None = None,
  30. style: Literal["%", "{", "$"] = "%",
  31. use_colors: bool | None = None,
  32. ):
  33. if use_colors in (True, False):
  34. self.use_colors = use_colors
  35. else:
  36. self.use_colors = sys.stdout.isatty()
  37. super().__init__(fmt=fmt, datefmt=datefmt, style=style)
  38. def color_level_name(self, level_name: str, level_no: int) -> str:
  39. def default(level_name: str) -> str:
  40. return str(level_name) # pragma: no cover
  41. func = self.level_name_colors.get(level_no, default)
  42. return func(level_name)
  43. def should_use_colors(self) -> bool:
  44. return True # pragma: no cover
  45. def formatMessage(self, record: logging.LogRecord) -> str:
  46. recordcopy = copy(record)
  47. levelname = recordcopy.levelname
  48. seperator = " " * (8 - len(recordcopy.levelname))
  49. if self.use_colors:
  50. levelname = self.color_level_name(levelname, recordcopy.levelno)
  51. if "color_message" in recordcopy.__dict__:
  52. recordcopy.msg = recordcopy.__dict__["color_message"]
  53. recordcopy.__dict__["message"] = recordcopy.getMessage()
  54. recordcopy.__dict__["levelprefix"] = levelname + ":" + seperator
  55. return super().formatMessage(recordcopy)
  56. class DefaultFormatter(ColourizedFormatter):
  57. def should_use_colors(self) -> bool:
  58. return sys.stderr.isatty() # pragma: no cover
  59. class AccessFormatter(ColourizedFormatter):
  60. status_code_colours = {
  61. 1: lambda code: click.style(str(code), fg="bright_white"),
  62. 2: lambda code: click.style(str(code), fg="green"),
  63. 3: lambda code: click.style(str(code), fg="yellow"),
  64. 4: lambda code: click.style(str(code), fg="red"),
  65. 5: lambda code: click.style(str(code), fg="bright_red"),
  66. }
  67. def get_status_code(self, status_code: int) -> str:
  68. try:
  69. status_phrase = http.HTTPStatus(status_code).phrase
  70. except ValueError:
  71. status_phrase = ""
  72. status_and_phrase = "%s %s" % (status_code, status_phrase)
  73. if self.use_colors:
  74. def default(code: int) -> str:
  75. return status_and_phrase # pragma: no cover
  76. func = self.status_code_colours.get(status_code // 100, default)
  77. return func(status_and_phrase)
  78. return status_and_phrase
  79. def formatMessage(self, record: logging.LogRecord) -> str:
  80. recordcopy = copy(record)
  81. (
  82. client_addr,
  83. method,
  84. full_path,
  85. http_version,
  86. status_code,
  87. ) = recordcopy.args # type: ignore[misc]
  88. status_code = self.get_status_code(int(status_code)) # type: ignore[arg-type]
  89. request_line = "%s %s HTTP/%s" % (method, full_path, http_version)
  90. if self.use_colors:
  91. request_line = click.style(request_line, bold=True)
  92. recordcopy.__dict__.update(
  93. {
  94. "client_addr": client_addr,
  95. "request_line": request_line,
  96. "status_code": status_code,
  97. }
  98. )
  99. return super().formatMessage(recordcopy)