Pasted

"""
File: logging.py
Author: Justin Garofolo
Date: 2024-05-31
Version: 1.1.0
 
Description:
    A flexible logging system for Python applications.
    It supports multiple logging levels and asynchronous log handling through threaded workers.
    The module can be easily integrated into projects with the capability to log messages both locally and remotely.
    Note that remote logging capabilities are implemented via BetterStack's logging service by default: https://betterstack.com/logs
 
Copyright:
    Copyright (C) 2024 SneakSync. All rights reserved.
 
Documentation:
    https://link.sneaksync.com/logging
"""
import requests, os, threading
from typing import Optional
from enum import Enum
from datetime import datetime, timezone
from http import HTTPStatus
from dataclasses import dataclass
from queue import Queue
 
def create_logger_from_env(
    name: str,
    env_var: str,
    auto_print: bool = True
) -> "Logger":
    tk = os.getenv(env_var)
    logger = Logger(name, token = tk, auto_print = auto_print)
    if not tk:
        logger.warn(f"Failed to initialize instance of Logger using create_logger_from_env. Missing environment variable: {env_var}")
    return logger
 
def create_logger(
    name: str,
    token: Optional[str] = None,
    auto_print: bool = True
) -> "Logger":
    return Logger(name, token = token, auto_print = auto_print)
 
class Level(Enum):
    TRACE = "trace"
    DEBUG = "debug"
    INFO = "info"
    WARN = "warn"
    ERROR = "error"
    FATAL = "fatal"
    PANIC = "panic"
 
@dataclass
class Log:
    id: int
    timestamp: datetime
    message: str
    level: Level
 
class Logger:
    _inc: int = 0 # auto-increment for unique log ID
 
    def __init__(
        self,
        name: str,
        token: Optional[str] = None,
        worker_count: int = 3,
        auto_print: bool = True
    ):
        self.name = name
        self.token = token
        self.worker_count = worker_count
        self.auto_print = auto_print
        self.queue = Queue()
    
        if token:
            mask = ('*' * (len(token) - 4) + token[-4:]) \
                if len(token) > 4 else token
            self.debug(f"Initialized logging module with token: {mask}")
 
    def log_worker(self, worker_id: int):
        self.info(f"Starting log worker (ID: {worker_id})")
        while True:
            log = self.queue.get()
 
            # thread kill switch
            if log is None:
                break
           
            if not isinstance(log, Log):
                type_str = type(log).__name__
                print(f"Item from log queue is not an instance of Log: {log} [{type_str}]")
                continue
 
            self.upload_log(log)
            self.queue.task_done()
        self.info(f"Killed log worker (ID: {worker_id})")
    
    def start_workers(self):
        for idx in range(self.worker_count):
            worker_id = idx + 1
            worker = threading.Thread(target = self.log_worker, args = (worker_id,))
            worker.daemon = True
            worker.start()
 
    def stop_workers(self):
        for _ in range(self.worker_count):
            self.queue.put(None)
 
    def log(self, level: Level, message: str):
        Logger._inc += 1
        log_id = Logger._inc
 
        if self.auto_print:
            now = datetime.now()
            now_str = now.strftime('%Y-%m-%d %I:%M:%S %p')
            print(f"[{self.name} | {now_str} | {level.value}] {message}")
 
        log_obj = Log(
            id = log_id,
            timestamp = datetime.now(timezone.utc),
            message = message,
            level = level
        )
        self.queue.put(log_obj)
 
    def debug(self, message: str):
        return self.log(Level.DEBUG, message)
 
    def info(self, message: str):
        return self.log(Level.INFO, message)
 
    def warn(self, message: str):
        return self.log(Level.WARN, message)
 
    def error(self, message: str):
        return self.log(Level.ERROR, message)
 
    def upload_log(self, log: Log) -> bool:
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f"Bearer {self.token}"
        }
        data = {
            "dt": log.timestamp.strftime('%Y-%m-%d %H:%M:%S.%f UTC'),
            # "message": f"[#{log.id}] {log.message}",
            "message": log.message,
            "level": log.level.value
        }
        url = "https://in.logs.betterstack.com"
        response = requests.post(url, json = data, headers = headers)
        if response.status_code != HTTPStatus.ACCEPTED.value:
            print(f"Failed to send log to logtail, status code: {response.status_code}")
            return False
        return True
 
__default_logger: Optional[Logger] = None
 
def set_default(logger: Logger):
    global __default_logger
    __default_logger = logger
 
def debug(message: str):
    assert __default_logger is not None, "Failed to invoke 'logging.debug': missing default logger."
    return __default_logger.debug(message)
 
def info(message: str):
    assert __default_logger is not None, "Failed to invoke 'logging.info': missing default logger."
    return __default_logger.info(message)
 
def warn(message: str):
    assert __default_logger is not None, "Failed to invoke 'logging.warn': missing default logger."
    return __default_logger.warn(message)
 
def error(message: str):
    assert __default_logger is not None, "Failed to invoke 'logging.error': missing default logger."
    return __default_logger.error(message)