Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 2,061 Bytes
38c5e59 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
import sys
import logging
import contextvars
from typing import Dict, Any
from .json_formatter import JsonFormatter
# Create a context variable to store request-specific information
context_var: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar(
"context_dict"
)
class Logger:
def __init__(self, context=None, use_context_var=False):
self.logger = logging.getLogger("json_logger")
self.context = context or {}
self.base_context = context or {}
self.use_context_var = use_context_var
self._setup()
def _setup(self):
self.logger.setLevel(logging.DEBUG)
self.context = (
context_var.set(self.base_context) if self.use_context_var else self.context
)
if not self.logger.handlers:
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setFormatter(JsonFormatter())
self.logger.addHandler(console_handler)
def debug(self, data, context={}):
self.log(logging.DEBUG, data, context)
def info(self, data, context={}):
self.log(logging.INFO, data, context)
def warning(self, data, context={}):
self.log(logging.WARNING, data, context)
def error(self, data, error=None, context={}):
self.log(logging.ERROR, data, context, error)
def log(self, level, data, context={}, error=None):
self.update_context(context=context)
self.logger.log(
level,
msg=data,
extra={"context": self._get_context()},
exc_info=error,
)
def _get_context(self):
return context_var.get() if self.use_context_var else self.context
def reset_context(self):
if self.use_context_var:
context_var.set(self.base_context)
else:
self.context = self.base_context or {}
def update_context(self, context):
if self.use_context_var:
context_var.set({**context_var.get(), **context})
else:
self.context.update(context)
|