nonebot2/nonebot/utils.py

import re
import json
import asyncio
import inspect
import collections
import dataclasses
from functools import wraps, partial
from contextlib import asynccontextmanager
from typing_extensions import GenericAlias  # type: ignore
from typing_extensions import ParamSpec, get_args, get_origin
from typing import (
    Any,
    Type,
    Deque,
    Tuple,
    Union,
    TypeVar,
    Callable,
    Optional,
    Awaitable,
    AsyncGenerator,
    ContextManager,
)

from nonebot.log import logger
from nonebot.typing import overrides

P = ParamSpec("P")
R = TypeVar("R")
T = TypeVar("T")


def escape_tag(s: str) -> str:
    """
    :说明:

      用于记录带颜色日志时转义 ``<tag>`` 类型特殊标签

    :参数:

      * ``s: str``: 需要转义的字符串

    :返回:

      - ``str``
    """
    return re.sub(r"</?((?:[fb]g\s)?[^<>\s]*)>", r"\\\g<0>", s)


def generic_check_issubclass(
    cls: Any, class_or_tuple: Union[Type[Any], Tuple[Type[Any], ...]]
) -> bool:
    try:
        return issubclass(cls, class_or_tuple)
    except TypeError:
        if get_origin(cls) is Union:
            for type_ in get_args(cls):
                if type_ is not type(None) and not generic_check_issubclass(
                    type_, class_or_tuple
                ):
                    return False
            return True
        elif isinstance(cls, GenericAlias):
            origin = get_origin(cls)
            return bool(origin and issubclass(origin, class_or_tuple))
        raise


def is_coroutine_callable(func: Callable[..., Any]) -> bool:
    if inspect.isroutine(func):
        return inspect.iscoroutinefunction(func)
    if inspect.isclass(func):
        return False
    func_ = getattr(func, "__call__", None)
    return inspect.iscoroutinefunction(func_)


def is_gen_callable(func: Callable[..., Any]) -> bool:
    if inspect.isgeneratorfunction(func):
        return True
    func_ = getattr(func, "__call__", None)
    return inspect.isgeneratorfunction(func_)


def is_async_gen_callable(func: Callable[..., Any]) -> bool:
    if inspect.isasyncgenfunction(func):
        return True
    func_ = getattr(func, "__call__", None)
    return inspect.isasyncgenfunction(func_)


def run_sync(func: Callable[P, R]) -> Callable[P, Awaitable[R]]:
    """
    :说明:

      一个用于包装 sync function 为 async function 的装饰器

    :参数:

      * ``func: Callable[P, R]``: 被装饰的同步函数

    :返回:

      - ``Callable[P, Awaitable[R]]``
    """

    @wraps(func)
    async def _wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        loop = asyncio.get_running_loop()
        pfunc = partial(func, *args, **kwargs)
        result = await loop.run_in_executor(None, pfunc)
        return result

    return _wrapper


@asynccontextmanager
async def run_sync_ctx_manager(
    cm: ContextManager[T],
) -> AsyncGenerator[T, None]:
    try:
        yield await run_sync(cm.__enter__)()
    except Exception as e:
        ok = await run_sync(cm.__exit__)(type(e), e, None)
        if not ok:
            raise e
    else:
        await run_sync(cm.__exit__)(None, None, None)


def get_name(obj: Any) -> str:
    if inspect.isfunction(obj) or inspect.isclass(obj):
        return obj.__name__
    return obj.__class__.__name__


class CacheLock:
    def __init__(self):
        self._waiters: Optional[Deque[asyncio.Future]] = None
        self._locked = False

    def __repr__(self):
        extra = "locked" if self._locked else "unlocked"
        if self._waiters:
            extra = f"{extra}, waiters: {len(self._waiters)}"
        return f"<{self.__class__.__name__} [{extra}]>"

    async def __aenter__(self):
        await self.acquire()
        return None

    async def __aexit__(self, exc_type, exc, tb):
        self.release()

    def locked(self):
        return self._locked

    async def acquire(self):
        if not self._locked and (
            self._waiters is None or all(w.cancelled() for w in self._waiters)
        ):
            self._locked = True
            return True

        if self._waiters is None:
            self._waiters = collections.deque()

        loop = asyncio.get_running_loop()
        future = loop.create_future()
        self._waiters.append(future)

        # Finally block should be called before the CancelledError
        # handling as we don't want CancelledError to call
        # _wake_up_first() and attempt to wake up itself.
        try:
            try:
                await future
            finally:
                self._waiters.remove(future)
        except asyncio.CancelledError:
            if not self._locked:
                self._wake_up_first()
            raise

        self._locked = True
        return True

    def release(self):
        if self._locked:
            self._locked = False
            self._wake_up_first()
        else:
            raise RuntimeError("Lock is not acquired.")

    def _wake_up_first(self):
        if not self._waiters:
            return
        try:
            future = next(iter(self._waiters))
        except StopIteration:
            return

        # .done() necessarily means that a waiter will wake up later on and
        # either take the lock, or, if it was cancelled and lock wasn't
        # taken already, will hit this again and wake up a new waiter.
        if not future.done():
            future.set_result(True)


class DataclassEncoder(json.JSONEncoder):
    """
    :说明:

      在JSON序列化 ``Message`` (List[Dataclass]) 时使用的 ``JSONEncoder``
    """

    @overrides(json.JSONEncoder)
    def default(self, o):
        if dataclasses.is_dataclass(o):
            return dataclasses.asdict(o)
        return super().default(o)


def logger_wrapper(logger_name: str):
    """
    :说明:

    用于打印 adapter 的日志。

    :log 参数:

    * ``level: Literal['WARNING', 'DEBUG', 'INFO']``: 日志等级
    * ``message: str``: 日志信息
    * ``exception: Optional[Exception]``: 异常信息
    """

    def log(level: str, message: str, exception: Optional[Exception] = None):
        return logger.opt(colors=True, exception=exception).log(
            level, f"<m>{escape_tag(logger_name)}</m> | " + message
        )

    return log