Python - Logging Module (logging.py)

"""
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)