nonebot-plugin-picstatus-re 2.3.0__py3-none-any.whl

nonebot_plugin_picstatus_re/__init__.py

@@ -0,0 +1,67 @@
+# ruff: noqa: E402
+
+from nonebot import get_driver, require
+from nonebot.plugin import PluginMetadata, inherit_supported_adapters
+
+require("nonebot_plugin_apscheduler")
+require("nonebot_plugin_alconna")
+require("nonebot_plugin_uninfo")
+require("nonebot_plugin_localstore")
+
+from . import __main__ as __main__, misc_statistics as misc_statistics
+from .bg_provider import bg_preloader
+from .collectors import (
+    enable_collectors,
+    load_builtin_collectors,
+    registered_collectors,
+)
+from .config import ConfigModel, config
+from .templates import load_builtin_templates, loaded_templates
+
+driver = get_driver()
+
+
+# lazy load builtin templates and collectors
+@driver.on_startup
+async def _():
+    if config.ps_template not in loaded_templates:
+        load_builtin_templates()
+    current_template = loaded_templates.get(config.ps_template)
+    if current_template is None:
+        raise ValueError(f"Template {config.ps_template} not found")
+
+    if (current_template.collectors is None) or any(
+        (x not in registered_collectors) for x in current_template.collectors
+    ):
+        load_builtin_collectors()
+
+    collectors = (
+        set(registered_collectors)
+        if current_template.collectors is None
+        else current_template.collectors
+    )
+    await enable_collectors(*collectors)
+
+    bg_preloader.start_preload()
+
+
+usage = f"指令：{' / '.join(config.ps_command)}"
+if config.ps_need_at:
+    usage += "\n注意：使用指令时需要@机器人"
+if config.ps_only_su:
+    usage += "\n注意：仅SuperUser可以使用此指令"
+
+__version__ = "2.3.0"
+__plugin_meta__ = PluginMetadata(
+    name="PicStatus-Re",
+    description="以图片形式显示当前设备的运行状态",
+    usage=usage,
+    type="application",
+    homepage="https://github.com/yuexps/nonebot-plugin-picstatus-re",
+    config=ConfigModel,
+    supported_adapters=inherit_supported_adapters(
+        "nonebot_plugin_alconna",
+        "nonebot_plugin_uninfo",
+    ),
+    extra={"License": "MIT", "Author": "yuexps"},
+)

nonebot_plugin_picstatus_re/__main__.py

@@ -0,0 +1,57 @@
+import asyncio
+
+from nonebot import logger, on_command
+from nonebot.adapters import Bot as BaseBot, Event as BaseEvent, Message as BaseMessage
+from nonebot.matcher import current_bot, current_event, current_matcher
+from nonebot.params import CommandArg
+from nonebot.permission import SUPERUSER
+from nonebot.rule import Rule, to_me
+from nonebot.typing import T_State
+from nonebot_plugin_alconna.uniseg import UniMessage
+
+from .bg_provider import bg_preloader
+from .collectors import collect_all
+from .config import config
+from .misc_statistics import bot_avatar_cache, bot_info_cache, cache_bot_avatar
+from .templates import render_current_template
+
+
+def check_empty_arg_rule(arg: BaseMessage = CommandArg()):
+    return not arg.extract_plain_text()
+
+
+def trigger_rule():
+    rule = Rule(check_empty_arg_rule)
+    if config.ps_need_at:
+        rule &= to_me()
+    return rule
+
+
+_cmd, *_alias = config.ps_command
+stat_matcher = on_command(
+    _cmd,
+    aliases=set(_alias),
+    rule=trigger_rule(),
+    permission=SUPERUSER if config.ps_only_su else None,
+)
+
+
+@stat_matcher.handle()
+async def _(bot: BaseBot, event: BaseEvent, state: T_State):
+    if (
+        (bot.self_id not in bot_avatar_cache)
+        and (info := bot_info_cache.get(bot.self_id))
+        and info.avatar
+    ):
+        await cache_bot_avatar(info.avatar, bot, event, state)
+
+    try:
+        bg, collected = await asyncio.gather(bg_preloader.get(), collect_all())
+        ret = await render_current_template(collected=collected, bg=bg)
+    except Exception:
+        logger.exception("获取运行状态图失败")
+        await UniMessage("获取运行状态图片失败，请检查后台输出").send(
+            reply_to=config.ps_reply_target,
+        )
+    else:
+        await UniMessage.image(raw=ret).send(reply_to=config.ps_reply_target)

nonebot_plugin_picstatus_re/bg_provider.py

@@ -0,0 +1,218 @@
+import asyncio as aio
+import mimetypes
+import random
+import sys
+import time
+from collections.abc import AsyncIterable
+from pathlib import Path
+from typing import NamedTuple, TypeAlias
+
+from cookit.common import race
+from cookit.loguru import warning_suppress
+from nonebot import get_driver, logger
+
+from .config import BG_PRELOAD_CACHE_DIR, DEFAULT_BG_PATH, ASSETS_PATH, config
+
+if sys.version_info >= (3, 11):
+    from asyncio.taskgroups import TaskGroup
+else:
+    from taskgroup import TaskGroup
+
+
+class BgBytesData(NamedTuple):
+    data: bytes | None
+    mime: str
+
+
+class BgFileData(NamedTuple):
+    path: Path | None
+    mime: str
+
+
+BgData: TypeAlias = BgBytesData | BgFileData
+DEFAULT_MIME = "application/octet-stream"
+
+
+def get_bg_files() -> list["Path"]:
+    if not config.ps_bg_local_path.exists():
+        logger.warning("Custom background path does not exist, fallback to default")
+        return [DEFAULT_BG_PATH]
+    if config.ps_bg_local_path.is_file():
+        return [config.ps_bg_local_path]
+
+    if config.ps_bg_local_path == ASSETS_PATH:
+        files = [x for x in config.ps_bg_local_path.glob("default_bg_*.webp") if x.is_file()]
+    else:
+        files = [
+            x for x in config.ps_bg_local_path.glob("*")
+            if x.is_file() and x.name != "default_avatar.webp"
+        ]
+    if not files:
+        logger.warning("Custom background dir has no file in it, fallback to default")
+        return [DEFAULT_BG_PATH]
+    return files
+
+
+BG_FILES = get_bg_files()
+
+
+def refresh_bg_files():
+    global BG_FILES
+    BG_FILES = get_bg_files()
+
+
+async def local(num: int) -> AsyncIterable[BgData]:
+    files = random.sample(BG_FILES, num)
+    for x in files:
+        yield BgFileData(
+            x,
+            mimetypes.guess_type(x)[0] or DEFAULT_MIME,
+        )
+
+
+async def none(num: int) -> AsyncIterable[BgData]:
+    for _ in range(num):
+        yield BgBytesData(None, DEFAULT_MIME)
+
+
+async def fetch_bg(num: int) -> AsyncIterable[BgData]:
+    provider = none if config.ps_bg_provider == "none" else local
+    async for x in provider(num):
+        yield x
+
+
+def cache_bg(bg: BgBytesData):
+    if not bg.data:
+        return BgFileData(None, bg.mime)
+    BG_PRELOAD_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    path = BG_PRELOAD_CACHE_DIR / f"{time.time_ns()}.{bg.mime.split('/')[-1]}"
+    path.write_bytes(bg.data)
+    return BgFileData(path, bg.mime)
+
+
+def read_cached_bg_file(bg: BgFileData) -> BgBytesData | None:
+    if not bg.path:
+        return BgBytesData(None, bg.mime)
+    with warning_suppress("Failed to read cached file"):
+        data = bg.path.read_bytes()
+        if bg.path.is_relative_to(BG_PRELOAD_CACHE_DIR):
+            with warning_suppress("Failed to unlink cached file"):
+                bg.path.unlink()
+        return BgBytesData(data, bg.mime)
+    return None
+
+
+async def get_one_fallback() -> BgBytesData:
+    with warning_suppress("Failed to get local bg file, fallback to none"):
+        async for x in local(1):
+            if bg := read_cached_bg_file(x):
+                return bg
+    logger.warning("Failed to read local bg file, fallback to none")
+    return BgBytesData(None, DEFAULT_MIME)
+
+
+class BgPreloader:
+    def __init__(self, preload_count: int):
+        self.preload_count = preload_count
+        self.background_queue = aio.Queue[BgData]()
+        self.current_load_task_main: aio.Task | None = None
+        self.consumed_in_loading: bool = False
+        self.image_got_signal = aio.Event()
+        self.fire_tasks: set[aio.Task] = set()
+
+    async def preload_task(
+        self,
+        count: int,
+        fire: bool = False,
+        fire_done_signal: aio.Event | None = None,
+    ):
+        logger.debug(f"Preload task started, will preload {count} images, {fire=}")
+        try:
+            async for x in fetch_bg(count):
+                logger.debug("Got one image")
+                if self.preload_count > 0 or (
+                    fire_done_signal and fire_done_signal.is_set()
+                ):
+                    x = cache_bg(x) if isinstance(x, BgBytesData) else x
+                await self.background_queue.put(x)
+                self.image_got_signal.set()
+                self.image_got_signal.clear()
+        except Exception:
+            logger.exception("Unexpected error occurred in preload task")
+        else:
+            logger.debug("Preload task finished")
+
+        if fire:
+            return
+        if (
+            self.consumed_in_loading
+            or self.background_queue.qsize() < self.preload_count
+        ):
+            self.consumed_in_loading = False
+            self.start_preload()
+        else:
+            self.current_load_task_main = None
+
+    def start_preload(self, force: bool = False):
+        count = self.preload_count - self.background_queue.qsize()
+        if count <= 0 and not force:
+            logger.debug(
+                "Current background queue size meets preload count, skip preload",
+            )
+            return
+        task = aio.create_task(self.preload_task(count))
+        self.current_load_task_main = task
+
+    def set_defer_preload(self):
+        if self.current_load_task_main:
+            logger.debug("Main preload task already running, set flag")
+            self.consumed_in_loading = True
+        else:
+            self.start_preload()
+
+    async def _get_on_fire(self) -> BgBytesData:
+        task_done_signal = aio.Event()
+        fire_task = aio.create_task(
+            self.preload_task(1, fire=True, fire_done_signal=task_done_signal),
+        )
+        fire_task.add_done_callback(lambda _: task_done_signal.set())
+        fire_task.add_done_callback(lambda _: self.fire_tasks.discard(fire_task))
+        self.fire_tasks.add(fire_task)
+        try:
+            await race(
+                task_done_signal.wait(),
+                aio.sleep(15),
+            )
+        finally:
+            task_done_signal.set()
+
+        if not self.background_queue.empty():
+            bg = await self.background_queue.get()
+            self.set_defer_preload()
+            if (not isinstance(bg, BgFileData)) or (bg := read_cached_bg_file(bg)):
+                return bg
+
+        logger.error("Unable to get an background image, falling back to local")
+        return await get_one_fallback()
+
+    async def get(self) -> BgBytesData:
+        self.set_defer_preload()
+
+        while not self.background_queue.empty():
+            bg = await self.background_queue.get()
+            self.set_defer_preload()
+            if (not isinstance(bg, BgFileData)) or (bg := read_cached_bg_file(bg)):
+                return bg
+
+        return await self._get_on_fire()
+
+
+bg_preloader = BgPreloader(config.ps_bg_preload_count)
+
+driver = get_driver()
+
+
+@driver.on_shutdown
+async def _():
+    for t in bg_preloader.fire_tasks:
+        t.cancel()

nonebot_plugin_picstatus_re/collectors/__init__.py

@@ -0,0 +1,226 @@
+import asyncio
+import importlib
+import time
+from abc import abstractmethod
+from collections import deque
+from collections.abc import Awaitable, Callable
+from contextlib import suppress
+from pathlib import Path
+from typing import Any, Generic, TypeVar
+from typing_extensions import override
+
+from nonebot import logger
+from nonebot_plugin_apscheduler import scheduler
+
+from ..config import config
+
+T = TypeVar("T")
+TI = TypeVar("TI")
+TR = TypeVar("TR")
+TC = TypeVar("TC", bound="Collector")
+TCF = TypeVar("TCF", bound=Callable[[], Awaitable[Any]])
+R = TypeVar("R")
+
+Undefined = type("Undefined", (), {})
+
+
+class SkipCollectError(Exception):
+    pass
+
+
+class Collector(Generic[TI, TR]):
+    @abstractmethod
+    async def _get(self) -> TI: ...
+
+    @abstractmethod
+    async def get(self) -> TR: ...
+
+
+class BaseNormalCollector(Collector[T, T], Generic[T]):
+    def __init__(self) -> None:
+        super().__init__()
+
+    @override
+    async def get(self) -> T:
+        return await self._get()
+
+
+class BaseFirstTimeCollector(Collector[T, T], Generic[T]):
+    def __init__(self) -> None:
+        super().__init__()
+        self._cached: T | Undefined = Undefined()
+
+    @override
+    async def get(self) -> T:
+        if not isinstance(self._cached, Undefined):
+            return self._cached
+        data = await self._get()
+        self._cached = data
+        return data
+
+
+class BasePeriodicCollector(Collector[T, deque[T]], Generic[T]):
+    def __init__(self, size: int = config.ps_default_collect_cache_size) -> None:
+        super().__init__()
+        self.data = deque(maxlen=size)
+
+    @override
+    async def get(self) -> deque[T]:
+        return self.data
+
+    async def collect(self):
+        try:
+            data = await self._get()
+        except SkipCollectError:
+            return
+        except Exception:
+            logger.exception("Error occurred while collecting data")
+        else:
+            self.data.append(data)
+
+
+registered_collectors: dict[str, type[Collector]] = {}
+enabled_collectors: dict[str, Collector] = {}
+
+
+def collector(name: str):
+    def deco(cls: type[TC]) -> type[TC]:
+        if name in registered_collectors:
+            raise ValueError(f"Collector {name} already exists")
+        registered_collectors[name] = cls
+        logger.debug(f"Registered collector {name}")
+        return cls
+
+    return deco
+
+
+def _enable_collector(name: str):
+    if name not in registered_collectors:
+        raise ValueError(f"Collector {name} not found")
+    cls = registered_collectors[name]
+    if issubclass(cls, BasePeriodicCollector) and name in config.ps_collect_cache_size:
+        instance = cls(size=config.ps_collect_cache_size[name])
+    else:
+        instance = cls()
+    enabled_collectors[name] = instance
+
+
+async def enable_collectors(*names: str):
+    for name in names:
+        _enable_collector(name)
+    await init_first_time_collectors()
+    await setup_periodic_collectors_update_job()
+
+
+def functional_collector(cls: type[Collector], name: str | None = None):
+    def deco(func: TCF) -> TCF:
+        collector_name = name or func.__name__
+        if not collector_name:
+            raise ValueError("name must be provided")
+
+        class Collector(cls):
+            async def _get(self) -> Any:
+                return await func()
+
+        collector(collector_name)(Collector)
+        return func
+
+    return deco
+
+
+def normal_collector(name: str | None = None):
+    return functional_collector(BaseNormalCollector, name)
+
+
+def first_time_collector(name: str | None = None):
+    return functional_collector(BaseFirstTimeCollector, name)
+
+
+def periodic_collector(name: str | None = None):
+    return functional_collector(BasePeriodicCollector, name)
+
+
+class BaseTimeBasedCounterCollector(Collector[R, Any], Generic[T, R]):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.last_obj: Undefined | T = Undefined()
+        self.last_time: float = 0
+        self.normal_delay: float = 1
+
+    @abstractmethod
+    async def _calc(self, past: T, now: T, time_passed: float) -> R: ...
+
+    @abstractmethod
+    async def _get_obj(self) -> T: ...
+
+    @override
+    async def _get(self) -> R:
+        past = self.last_obj
+        past_time = self.last_time
+        time_now = time.time()
+        time_passed = time_now - past_time
+
+        self.last_time = time_now
+        self.last_obj = await self._get_obj()
+        if isinstance(past, Undefined):
+            raise SkipCollectError
+        return await self._calc(past, self.last_obj, time_passed)
+
+
+class NormalTimeBasedCounterCollector(
+    BaseTimeBasedCounterCollector[T, R],
+    BaseNormalCollector[R],
+    Generic[T, R],
+):
+    @override
+    async def get(self) -> R:
+        with suppress(SkipCollectError):
+            await self._get()
+        await asyncio.sleep(self.normal_delay)
+        return await self._get()
+
+
+class PeriodicTimeBasedCounterCollector(
+    BaseTimeBasedCounterCollector[T, R],
+    BasePeriodicCollector[R],
+    Generic[T, R],
+): ...
+
+
+async def collect_all() -> dict[str, Any]:
+    async def get(name: str):
+        return name, await enabled_collectors[name].get()
+
+    res = await asyncio.gather(*(get(name) for name in enabled_collectors))
+    return dict(res)
+
+
+def load_builtin_collectors():
+    for module in Path(__file__).parent.iterdir():
+        if not module.name.startswith("_"):
+            importlib.import_module(f".{module.stem}", __package__)
+
+
+async def init_first_time_collectors():
+    await asyncio.gather(
+        *(
+            x.get()
+            for x in enabled_collectors.values()
+            if isinstance(x, BaseFirstTimeCollector)
+        ),
+    )
+
+
+async def setup_periodic_collectors_update_job():
+    collectors = [
+        x for x in enabled_collectors.values() if isinstance(x, BasePeriodicCollector)
+    ]
+    if not collectors:
+        return
+    logger.debug("Setting up periodic collectors")
+
+    @scheduler.scheduled_job("interval", seconds=config.ps_collect_interval)
+    async def _do():
+        await asyncio.gather(*(x.collect() for x in collectors))
+
+    await _do()

nonebot_plugin_picstatus_re/collectors/bot.py

@@ -0,0 +1,93 @@
+import asyncio
+from dataclasses import dataclass
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+from nonebot import get_bots, logger
+from nonebot.matcher import current_bot
+
+from ..config import config
+from ..misc_statistics import bot_connect_time, bot_info_cache, recv_num, send_num
+from ..util import format_timedelta
+from . import normal_collector
+
+if TYPE_CHECKING:
+    from nonebot.adapters import Bot as BaseBot
+
+try:
+    from nonebot.adapters.onebot.v11 import Bot as OBV11Bot
+except ImportError:
+    OBV11Bot = None
+
+
+@dataclass
+class BotStatus:
+    self_id: str
+    adapter: str
+    nick: str
+    bot_connected: str
+    msg_rec: str
+    msg_sent: str
+
+
+async def get_ob11_msg_num(bot: "BaseBot") -> tuple[int | None, int | None]:
+    if not (config.ps_ob_v11_use_get_status and OBV11Bot and isinstance(bot, OBV11Bot)):
+        return None, None
+
+    try:
+        bot_stat = (await bot.get_status()).get("stat")
+    except Exception as e:
+        logger.warning(
+            f"Error when getting bot status: {e.__class__.__name__}: {e}",
+        )
+        return None, None
+    if not bot_stat:
+        return None, None
+
+    msg_rec = bot_stat.get("message_received") or bot_stat.get(
+        "MessageReceived",
+    )
+    msg_sent = bot_stat.get("message_sent") or bot_stat.get("MessageSent")
+    return msg_rec, msg_sent
+
+
+async def get_bot_status(bot: "BaseBot", now_time: datetime) -> BotStatus:
+    nick = (
+        ((info := bot_info_cache[bot.self_id]).nick or info.name or info.id)
+        if (not config.ps_use_env_nick) and (bot.self_id in bot_info_cache)
+        else next(iter(config.nickname), None)
+    ) or "Bot"
+    bot_connected = (
+        format_timedelta(now_time - t)
+        if (t := bot_connect_time.get(bot.self_id))
+        else "未知"
+    )
+
+    msg_rec, msg_sent = await get_ob11_msg_num(bot)
+    if msg_rec is None:
+        msg_rec = recv_num.get(bot.self_id)
+    if msg_sent is None:
+        msg_sent = send_num.get(bot.self_id)
+    msg_rec = "未知" if (msg_rec is None) else str(msg_rec)
+    msg_sent = "未知" if (msg_sent is None) else str(msg_sent)
+
+    return BotStatus(
+        self_id=bot.self_id,
+        adapter=bot.adapter.get_name(),
+        nick=nick,
+        bot_connected=bot_connected,
+        msg_rec=msg_rec,
+        msg_sent=msg_sent,
+    )
+
+
+@normal_collector()
+async def bots() -> list[BotStatus]:
+    now_time = datetime.now().astimezone()
+    return (
+        [await get_bot_status(current_bot.get(), now_time)]
+        if config.ps_show_current_bot_only
+        else await asyncio.gather(
+            *(get_bot_status(bot, now_time) for bot in get_bots().values()),
+        )
+    )