Source code for aiomisc.log

import asyncio
import atexit
import logging
import logging.handlers
import time
import traceback
from contextlib import suppress
from functools import partial
from socket import socket
from typing import (
    Any, Callable, Dict, Iterable, List, Optional, Tuple, Type, Union,
)
from weakref import finalize

import aiomisc_log
from aiomisc_log.enum import LogFormat, LogLevel

from .thread_pool import run_in_new_thread


def _thread_flusher(
    handler: logging.handlers.MemoryHandler,
    flush_interval: Union[float, int],
    loop: asyncio.AbstractEventLoop,
) -> None:
    def has_no_target() -> bool:
        return True

    def has_target() -> bool:
        return bool(handler.target)

    is_target = has_no_target

    if isinstance(handler, logging.handlers.MemoryHandler):
        is_target = has_target

    while not loop.is_closed() and is_target():
        with suppress(Exception):
            if handler.buffer:
                handler.flush()

        time.sleep(flush_interval)


def suppressor(
    callback: Callable[..., None],
    exceptions: Tuple[Type[BaseException], ...] = (Exception,),
) -> Callable[..., None]:
    def wrapper() -> None:
        with suppress(*exceptions):
            callback()
    return wrapper


def wrap_logging_handler(
    handler: logging.Handler,
    loop: Optional[asyncio.AbstractEventLoop] = None,
    buffer_size: int = 1024,
    flush_interval: Union[float, int] = 0.1,
) -> logging.Handler:
    buffered_handler = logging.handlers.MemoryHandler(
        buffer_size,
        target=handler,
        flushLevel=logging.CRITICAL,
    )

    run_in_new_thread(
        _thread_flusher, args=(
            buffered_handler, flush_interval, loop,
        ), no_return=True, statistic_name="logger",
    )

    at_exit_flusher = suppressor(handler.flush)
    atexit.register(at_exit_flusher)
    finalize(buffered_handler, partial(atexit.unregister, at_exit_flusher))

    return buffered_handler


class UnhandledLoopHook(aiomisc_log.UnhandledHook):
    @staticmethod
    def _fill_transport_extra(
        transport: Optional[asyncio.Transport],
        extra: Dict[str, Any],
    ) -> None:
        if transport is None:
            return

        extra["transport"] = repr(transport)

        for key in (
            "peername", "sockname", "compression",
            "cipher", "peercert", "pipe", "subprocess",
        ):
            value = transport.get_extra_info(key)
            if value:
                extra[f"transport_{key}"] = value

    def __call__(
        self, loop: asyncio.AbstractEventLoop,
        context: Dict[str, Any],
    ) -> None:
        context = dict(context)
        message: str = context.pop("message", "unhandled loop exception")
        exception: Optional[BaseException] = context.pop("exception", None)
        future: Optional[asyncio.Future] = context.pop("future", None)
        task: Optional[asyncio.Task] = context.pop("task", None)
        handle: Optional[asyncio.Handle] = context.pop("handle", None)
        protocol: Optional[asyncio.Protocol] = context.pop("protocol", None)
        transport: Optional[asyncio.Transport] = context.pop("transport", None)
        sock: Optional[socket] = context.pop("socket", None)
        source_traceback: List[traceback.FrameSummary] = (
            context.pop("source_traceback", None) or []
        )

        if exception is None:
            if future is not None:
                exception = future.exception()
            elif task is not None and task.done():
                exception = task.exception()

        extra = context
        if handle is not None:
            extra["handle"] = repr(handle)
        if protocol is not None:
            extra["protocol"] = repr(protocol)
        if sock is not None:
            extra["sock"] = repr(sock)

        self._fill_transport_extra(transport, extra)
        self.logger.exception(message, exc_info=exception, extra=extra)
        if source_traceback:
            self.logger.error(
                "".join(traceback.format_list(source_traceback)),
            )


def basic_config(
    level: Union[int, str] = LogLevel.default(),
    log_format: Union[str, LogFormat] = LogFormat.default(),
    buffered: bool = True, buffer_size: int = 1024,
    flush_interval: Union[int, float] = 0.2,
    loop: Optional[asyncio.AbstractEventLoop] = None,
    handlers: Iterable[logging.Handler] = (),
    **kwargs: Any,
) -> None:
    loop = loop or asyncio.get_event_loop()
    unhandled_hook = UnhandledLoopHook(logger_name="asyncio.unhandled")

    def wrap_handler(handler: logging.Handler) -> logging.Handler:
        nonlocal buffer_size, buffered, loop, unhandled_hook

        unhandled_hook.add_handler(handler)

        if buffered:
            return wrap_logging_handler(
                handler=handler,
                buffer_size=buffer_size,
                flush_interval=flush_interval,
                loop=loop,
            )
        return handler

    aiomisc_log.basic_config(
        level=level,
        log_format=log_format,
        handler_wrapper=wrap_handler,
        handlers=handlers,
        **kwargs,
    )

    loop.set_exception_handler(unhandled_hook)


__all__ = (
    "LogFormat",
    "LogLevel",
    "basic_config",
)