moemoe 0.0.1__py3-none-any.whl

moemoe/__init__.py

@@ -0,0 +1,60 @@
+"""moemoe —— QQ 官方机器人的应用端框架，配合 robot 协议端。
+
+    QQ 官方 ←WSS→ robot（协议端）←反向 WS→ moemoe（应用端）
+
+robot 主动拨号，所以 moemoe 是 WS 服务端。一个 Moemoe 同时接多个 robot，每个 robot
+一个 `Bot` 实例；`Event.bot` 指回收到这条消息的那个。
+
+    app / bot                  应用对象与已连接的机器人
+    event / message            事件与消息，纯数据
+    command / permission       分发
+    loader                     扫包加载插件
+    errors                     异常
+
+三方依赖只有 starlette 和 loguru。`from moemoe import logger` 拿到的就是 loguru 的
+全局 logger 原样再导出，`logger.disable('moemoe')` 让库彻底闭嘴。
+
+设计理由见 docs/design.md，使用者要知道的行为和限制见 README.md。
+"""
+
+from loguru import logger
+
+from .app import MESSAGE_EVENTS, Failure, Moemoe
+from .bot import Bot
+from .command import Args, Command, Registry
+from .errors import CallError, SessionTimeout, TooManySessions
+from .event import Event
+from .loader import load_plugins
+from .message import At, Image, Keyboard, Markdown, Message, Segment, Text, coerce
+from .session import Session, Sessions
+from .permission import (
+    EVERYBODY,
+    GROUP,
+    GROUP_ADMIN,
+    GROUP_OWNER,
+    NOBODY,
+    PRIVATE,
+    SUPERUSER,
+    Permission,
+)
+
+__version__ = '0.0.1'
+
+__all__ = [
+    # 应用
+    'Moemoe', 'Bot', 'Event', 'Failure', 'load_plugins', 'MESSAGE_EVENTS',
+    # 命令
+    'Args', 'Command', 'Registry',
+    # 会话
+    'Session', 'Sessions',
+    # 消息
+    'Message', 'Segment', 'Text', 'At', 'Image', 'Markdown', 'Keyboard', 'coerce',
+    # 权限
+    'Permission', 'EVERYBODY', 'NOBODY', 'GROUP', 'PRIVATE',
+    'GROUP_ADMIN', 'GROUP_OWNER', 'SUPERUSER',
+    # 异常
+    'CallError', 'SessionTimeout', 'TooManySessions',
+    # 日志，loguru 的全局 logger 原样再导出
+    'logger',
+    '__version__',
+]

moemoe/app.py

@@ -0,0 +1,412 @@
+"""Moemoe —— 应用对象：接连接、分发事件、跑命令。
+
+robot 主动拨号，所以这边是 WS 服务端。一个 Moemoe 同时接多个 robot，每个一个 Bot。
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
+from contextlib import asynccontextmanager
+from typing import Any
+
+from loguru import logger
+from starlette.applications import Starlette
+from starlette.routing import Route, WebSocketRoute
+from starlette.websockets import WebSocket, WebSocketDisconnect
+
+from .bot import Bot
+from .command import Command, Handler as CommandHandler, Registry
+from .errors import CallError
+from .event import Event
+from .session import Sessions
+
+__all__ = ['Moemoe', 'Failure', 'MESSAGE_EVENTS']
+
+Handler = Callable[..., Awaitable[Any]]
+
+# 消息类事件。`on_message` 订的就是这三种，命令也只在这三种上跑
+MESSAGE_EVENTS = (
+    'GROUP_AT_MESSAGE_CREATE',
+    'GROUP_MESSAGE_CREATE',
+    'C2C_MESSAGE_CREATE',
+)
+
+
+class Failure:
+    """一次处理函数出错，交给 `on_error` 钩子
+
+    :ivar error: 抛出来的异常
+    :ivar handler: 出错的处理函数
+    :ivar args: 当时传给它的参数，消息类事件时第一个是 `Event`
+    """
+
+    def __init__(self, error: BaseException, handler: Handler,
+                 args: tuple[Any, ...]) -> None:
+        self.error = error
+        self.handler = handler
+        self.args = args
+
+    def __repr__(self) -> str:
+        name = getattr(self.handler, '__name__', repr(self.handler))
+        return f'<Failure {name} {type(self.error).__name__}>'
+
+    @property
+    def event(self) -> Event | None:
+        """出错时手上那条事件，钩子不是事件触发的就是 None"""
+        first = self.args[0] if self.args else None
+        return first if isinstance(first, Event) else None
+
+
+class Moemoe:
+    """应用对象
+
+    :param access_token: 和 robot 配的一致，留空表示不校验
+    :param path: WebSocket 路径，robot 的 `backend.url` 指到这儿
+    :param timeout: 调用等回执多久算超时，秒
+    :param prefixes: 命令前缀，`('/', '')` 表示带不带斜杠都认
+    :param nicknames: 昵称触发，「萌萌，查天气」也能触发命令
+    :param superusers: 超级用户的 openid，`SUPERUSER` 谓词读它
+    :param sessions: 最多同时挂多少个会话
+    """
+
+    def __init__(self, *, access_token: str = '', path: str = '/robot',
+                 timeout: float = 30.0,
+                 prefixes: Sequence[str] = ('/',),
+                 nicknames: Sequence[str] = (),
+                 superusers: Sequence[str] = (),
+                 sessions: int = 1000) -> None:
+        self.access_token = access_token
+        self.timeout = timeout
+        self.superusers = tuple(superusers)
+
+        self.commands = Registry(prefixes, nicknames=nicknames)
+        self.bots: dict[str, Bot] = {}
+        self.sessions = Sessions(sessions)
+
+        self._handlers: list[tuple[str | None, Handler]] = []
+        self._connect: list[Handler] = []
+        self._disconnect: list[Handler] = []
+        self._startup: list[Handler] = []
+        self._shutdown: list[Handler] = []
+        self._errors: list[Handler] = []
+        # 强引用住在跑的任务，事件循环只持弱引用
+        self._tasks: set[asyncio.Task] = set()
+
+        self.asgi = Starlette(routes=[WebSocketRoute(path, self._serve)],
+                              lifespan=self._lifespan)
+
+    def __repr__(self) -> str:
+        return f'<Moemoe {len(self.bots)} bots, {len(self.commands)} commands>'
+
+    def bot(self, id: str | None = None) -> Bot:
+        """拿一个已连上的机器人
+
+        :param id: 机器人的 app_id，不给且只连了一个时返回那一个
+        :return: Bot
+        :raises CallError: 没有已连上的机器人，或者指名的那个不在
+        """
+        if id is not None:
+            if found := self.bots.get(id):
+                return found
+            raise CallError(0, {'error': f'机器人 {id} 没连上'})
+
+        if len(self.bots) == 1:
+            return next(iter(self.bots.values()))
+        if not self.bots:
+            raise CallError(0, {'error': '还没有机器人连上来'})
+        raise CallError(0, {'error': f'连着 {len(self.bots)} 个机器人，要指名是哪个'})
+
+    # ── 注册 ──
+
+    def command(self, name: str, **kwargs: Any) -> Callable[[CommandHandler], CommandHandler]:
+        """注册命令，处理函数签名 `(event, args)`，返回非 None 会当消息发出去
+
+        :param name: 命令名
+        :param kwargs: alias / patterns / permission / denied / to_me，见 `Command`
+        :return: 装饰器
+        """
+
+        def wrap(func: CommandHandler) -> CommandHandler:
+            self.commands.add(Command(name, func, **kwargs))
+            return func
+
+        return wrap
+
+    def on_event(self, kind: str | None = None) -> Callable[[Handler], Handler]:
+        """注册事件处理函数，签名 `(event)`
+
+        :param kind: 事件类型，不给表示全收
+        :return: 装饰器
+        """
+
+        def wrap(func: Handler) -> Handler:
+            self._handlers.append((kind, func))
+            return func
+
+        return wrap
+
+    def on_message(self) -> Callable[[Handler], Handler]:
+        """注册消息处理函数，签名 `(event)`，订 `MESSAGE_EVENTS` 那三种
+
+        :return: 装饰器
+        """
+
+        def wrap(func: Handler) -> Handler:
+            for kind in MESSAGE_EVENTS:
+                self._handlers.append((kind, func))
+            return func
+
+        return wrap
+
+    def on_bot_connect(self, func: Handler) -> Handler:
+        """注册连接建立钩子，签名 `(bot)`，钩子里可以调接口
+
+        :param func: 钩子
+        :return: 原样返回
+        """
+        self._connect.append(func)
+        return func
+
+    def on_bot_disconnect(self, func: Handler) -> Handler:
+        """注册连接断开钩子，签名 `(bot)`
+
+        :param func: 钩子
+        :return: 原样返回
+        """
+        self._disconnect.append(func)
+        return func
+
+    def on_startup(self, func: Handler) -> Handler:
+        """注册启动钩子，签名 `()`，在事件循环里跑
+
+        :param func: 钩子
+        :return: 原样返回
+        """
+        self._startup.append(func)
+        return func
+
+    def on_shutdown(self, func: Handler) -> Handler:
+        """注册关闭钩子，签名 `()`
+
+        :param func: 钩子
+        :return: 原样返回
+        """
+        self._shutdown.append(func)
+        return func
+
+    def on_error(self, func: Handler) -> Handler:
+        """注册出错钩子，签名 `(error)`，参数是 `Failure`
+
+        钩子自己抛异常只记一行，不再往上抛。
+
+        :param func: 钩子
+        :return: 原样返回
+        """
+        self._errors.append(func)
+        return func
+
+    def route(self, path: str, methods: Sequence[str] = ('GET',),
+              **kwargs: Any) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        """注册一条 HTTP 路由，处理函数签名 `(request)`
+
+        挂在同一个 ASGI 应用上，跟 WebSocket 那条路共用端口。
+
+        :param path: 路径
+        :param methods: HTTP 方法
+        :param kwargs: 透传给 `starlette.routing.Route`
+        :return: 装饰器
+        """
+
+        def wrap(func: Callable[..., Any]) -> Callable[..., Any]:
+            self.asgi.routes.append(Route(path, func, methods=list(methods), **kwargs))
+            return func
+
+        return wrap
+
+    # ── 生命周期 ──
+
+    @asynccontextmanager
+    async def _lifespan(self, _app: Starlette) -> AsyncIterator[None]:
+        """ASGI 的启动和关闭"""
+        await self._boot()
+        try:
+            yield
+        finally:
+            await self._halt()
+
+    async def _boot(self) -> None:
+        """跑启动钩子"""
+        for hook in self._startup:
+            await self._guard(hook)
+
+    async def _halt(self) -> None:
+        """关掉所有会话，跑关闭钩子"""
+        self.sessions.close()
+        for hook in self._shutdown:
+            await self._guard(hook)
+
+    # ── 连接 ──
+
+    async def _serve(self, ws: WebSocket) -> None:
+        """一条 robot 连接的生命周期"""
+        token = (ws.headers.get('authorization') or '').removeprefix('Bearer ').strip()
+        if self.access_token and token != self.access_token:
+            logger.warning('access_token 不对，拒了来自 {} 的连接', ws.client)
+            await ws.close(code=1008)
+            return
+
+        id = ws.headers.get('x-bot-id') or ''
+        if not id:
+            logger.warning('连接没带 X-Bot-Id，拒了')
+            await ws.close(code=1008)
+            return
+
+        await ws.accept()
+        bot = Bot(id, ws, app=self, timeout=self.timeout)
+        if old := self.bots.get(id):
+            old.abort(f'机器人 {id} 换了一条新连接')
+        self.bots[id] = bot
+        logger.info('{} 连上了，现在连着 {} 个机器人', bot, len(self.bots))
+
+        for hook in self._connect:
+            self._spawn(hook, bot, bot=bot)
+
+        try:
+            while True:
+                await self._handle(bot, json.loads(await ws.receive_text()))
+        except WebSocketDisconnect:
+            logger.warning('{} 断开', bot)
+        except Exception as error:  # noqa: BLE001
+            logger.exception('{} 连接异常: {}', bot, error)
+        finally:
+            if self.bots.get(id) is bot:
+                del self.bots[id]
+            bot.abort('robot 断开了')
+            # 断开钩子不归这个机器人管 —— abort 刚把它的任务全取消了
+            for hook in self._disconnect:
+                self._spawn(hook, bot)
+
+    async def _handle(self, bot: Bot, frame: dict) -> None:
+        """处理一帧"""
+        kind = frame.get('type')
+
+        if kind == 'result':
+            bot.settle(str(frame.get('echo') or ''), frame)
+            return
+
+        if kind == 'robot':
+            self._status(bot, frame)
+            return
+
+        if kind == 'event':
+            event = Event(frame.get('payload') or {}, bot)
+
+            # 会话优先。答案同时触发一条命令的话，多轮问答就没法用了
+            if event.type in MESSAGE_EVENTS and self.sessions.deliver(event):
+                return
+
+            if event.type in MESSAGE_EVENTS and len(self.commands):
+                self._spawn(self._dispatch, event, bot=bot)
+            for want, handler in self._handlers:
+                if want is None or want == event.type:
+                    self._spawn(handler, event, bot=bot)
+            return
+
+        logger.debug('{} 发来不认识的帧: {}', bot, kind)
+
+    def _status(self, bot: Bot, frame: dict) -> None:
+        """robot 自己的状态帧"""
+        data = frame.get('data') or {}
+        event = frame.get('event')
+
+        if event == 'ready':
+            bot.openid = data.get('openid') or ''
+            bot.name = data.get('name') or ''
+            logger.info('{} 就绪，openid={}', bot, bot.openid or '(robot 还没问到)')
+        elif event == 'gateway' and not data.get('up'):
+            logger.warning('{} 的 QQ 网关断了: {}', bot, data.get('reason', ''))
+
+    # ── 分发 ──
+
+    async def _dispatch(self, event: Event) -> None:
+        """认出命令就跑，返回非 None 当消息发出去"""
+        found = self.commands.parse(event.text)
+        if found is None:
+            logger.debug('{} 认不出命令: {}', event.bot, event.text)
+            return
+        command, args = found
+
+        if command.to_me and event.is_group and not event.to_me:
+            logger.debug('{} 不是冲机器人来的，挡了 {}', event.user_openid, command.name)
+            return
+
+        if command.permission is not None and not command.permission(event, self.superusers):
+            logger.debug('{} 权限不够，挡了 {}', event.user_openid, command.name)
+            if command.denied:
+                await event.send(command.denied)
+            return
+
+        out = await command.handler(event, args)
+        if out is not None:
+            await event.send(out)
+
+    def _spawn(self, handler: Handler, *args: Any, bot: Bot | None = None) -> None:
+        """起一个任务跑处理函数，异常在任务里就地记掉
+
+        所有处理函数走同一个错误边界：命令、事件处理函数、各种钩子。
+
+        给了 bot 就交给它管，那条连接断开时一起取消。
+
+        :param handler: 处理函数
+        :param args: 传给它的参数
+        :param bot: 这个处理函数属于哪个机器人
+        """
+        task = asyncio.create_task(self._guard(handler, *args))
+        if bot is not None:
+            bot.hold(task)
+            return
+
+        # 不属于任何机器人的（启动、关闭钩子），应用自己留一份强引用
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.discard)
+
+    async def _guard(self, handler: Handler, *args: Any) -> None:
+        """跑一个处理函数，出错只记不抛，并交给 on_error 钩子"""
+        name = getattr(handler, '__name__', repr(handler))
+        try:
+            await handler(*args)
+        except CallError as error:
+            logger.error('{} 调用失败: {}', name, error)
+            await self._report(Failure(error, handler, args))
+        except Exception as error:  # noqa: BLE001
+            logger.exception('{} 出错: {}', name, error)
+            await self._report(Failure(error, handler, args))
+
+    async def _report(self, failure: Failure) -> None:
+        """把一次出错交给 on_error 钩子，钩子自己抛了只记一行"""
+        for hook in self._errors:
+            try:
+                await hook(failure)
+            except Exception as error:  # noqa: BLE001
+                logger.exception('on_error 钩子自己出错了: {}', error)
+
+    # ── 跑起来 ──
+
+    def run(self, host: str = '127.0.0.1', port: int = 8000, **kwargs: Any) -> None:
+        """起服务等 robot 连进来
+
+        挂进已有的 ASGI 应用就别用这个，把 `asgi` 交给你的服务器。
+
+        :param host: 监听地址
+        :param port: 监听端口
+        :param kwargs: 透传给 uvicorn
+        """
+        try:
+            import uvicorn
+        except ImportError:
+            raise RuntimeError('没装 uvicorn：pip install moemoe[uvicorn]，'
+                               '或者把 Moemoe.asgi 交给你自己的 ASGI 服务器') from None
+
+        uvicorn.run(self.asgi, host=host, port=port, **kwargs)

moemoe/bot.py

@@ -0,0 +1,180 @@
+"""Bot —— 一个已连上的 robot，代表一个 QQ 机器人。
+
+一个 Moemoe 同时接多个 robot，每个 robot 一个 Bot 实例。发消息、调接口都经过这里，
+所以走的一定是这个机器人自己那条连接。
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+from loguru import logger
+
+from .errors import CallError
+from .message import Message, Segment, coerce
+
+if TYPE_CHECKING:
+    from starlette.websockets import WebSocket
+
+__all__ = ['Bot']
+
+
+class Bot:
+    """一个已连上的 robot
+
+    :param id: 机器人的 app_id，来自连接时的 `X-Bot-Id` 头
+    :param ws: 这个机器人的 WebSocket 连接
+    :param app: 收留它的应用对象
+    :param timeout: 等回执多久算超时，秒
+    """
+
+    def __init__(self, id: str, ws: WebSocket, *, app: Any = None,
+                 timeout: float = 30.0) -> None:
+        self.id = id
+        self.app = app
+        self.timeout = timeout
+        # 机器人自己的 openid 和显示名，robot 的 ready 帧带过来
+        self.openid: str = ''
+        self.name: str = ''
+
+        self._ws = ws
+        self._waiting: dict[str, asyncio.Future] = {}
+        self._echo = 0
+        # 这个机器人的处理函数。断开时一起取消，而且事件循环只持弱引用，
+        # 不自己留一份的话任务可能跑到一半被回收
+        self._tasks: set[asyncio.Task] = set()
+
+    def __repr__(self) -> str:
+        return f'<Bot {self.name or self.id}>'
+
+    async def call(self, method: str, path: str, body: dict | None = None,
+                   *, borrow: bool = False) -> dict:
+        """调一次官方接口
+
+        :param method: HTTP 方法
+        :param path: 接口路径，照腾讯的文档写。`/robot/` 开头的是 robot 自己的接口
+        :param body: 请求体，字段照腾讯的文档写
+        :param borrow: 让 robot 补一条这个会话还没过期的被动回复凭据
+        :return: QQ 的响应体
+        :raises CallError: 连接不在、超时，或者 QQ 回了 4xx / 5xx
+        """
+        self._echo += 1
+        echo = str(self._echo)
+        future: asyncio.Future = asyncio.get_running_loop().create_future()
+        self._waiting[echo] = future
+
+        frame: dict[str, Any] = {'type': 'call', 'echo': echo,
+                                 'method': method, 'path': path}
+        if body is not None:
+            frame['body'] = body
+        if borrow:
+            frame['passive'] = 'auto'
+
+        try:
+            await self._ws.send_text(json.dumps(frame, ensure_ascii=False))
+            result = await asyncio.wait_for(future, self.timeout)
+        except asyncio.TimeoutError:
+            raise CallError(0, {'error': f'{method} {path} 等回执超过 {self.timeout} 秒'}) from None
+        finally:
+            self._waiting.pop(echo, None)
+
+        status = result.get('status') or 0
+        payload = result.get('body')
+        if status < 200 or status >= 300:
+            raise CallError(status, payload, result.get('trace') or '')
+        return payload if isinstance(payload, dict) else {}
+
+    async def upload(self, kind: str, target: str, url: str) -> str:
+        """把一个公网图片交给 QQ 转存
+
+        :param kind: `groups` 或 `users`
+        :param target: 群 openid 或单聊对方的 openid
+        :param url: 公网可访问的图片地址
+        :return: 发消息要用的 `file_info`
+        """
+        out = await self.call('POST', f'/v2/{kind}/{target}/files',
+                              {'file_type': 1, 'url': url, 'srv_send_msg': False})
+        return out.get('file_info') or ''
+
+    async def deliver(self, kind: str, target: str,
+                      message: Message | Segment | str | dict,
+                      extra: dict[str, Any],
+                      next_seq: Callable[[], int] | None,
+                      *, borrow: bool = False) -> list[dict]:
+        """把一条消息发出去，必要时拆成几条
+
+        :param kind: `groups` 或 `users`
+        :param target: 群 openid 或单聊对方的 openid
+        :param message: 消息段、字符串，或者当请求体用的 dict
+        :param extra: 补进每个请求体的字段，被动回复凭据走这儿
+        :param next_seq: 每条消息的 `msg_seq`，主动消息传 None
+        :param borrow: 让 robot 补一条还没过期的凭据
+        :return: 每条消息的响应体
+        """
+        built = coerce(message)
+        bodies = [built] if isinstance(built, dict) else await built.build(self, kind, target)
+
+        sent = []
+        for body in bodies:
+            payload = {**body, **extra}
+            if next_seq is not None:
+                payload['msg_seq'] = next_seq()
+            sent.append(await self.call('POST', f'/v2/{kind}/{target}/messages',
+                                        payload, borrow=borrow))
+        return sent
+
+    async def push(self, target: str, message: Message | Segment | str | dict,
+                   *, private: bool = False, borrow: bool = False) -> list[dict]:
+        """主动推送到某个会话
+
+        不带 `msg_id`，也不带 `message_reference`。
+
+        :param target: 群 openid，或单聊对方的 openid
+        :param message: 消息段、字符串，或者当请求体用的 dict
+        :param private: 发单聊
+        :param borrow: 让 robot 找一条这个会话还没过期的凭据，找到就免费
+        :return: 每条消息的响应体
+        """
+        kind = 'users' if private else 'groups'
+        return await self.deliver(kind, target, message, {}, None, borrow=borrow)
+
+    def settle(self, echo: str, result: dict) -> None:
+        """回执到了，唤醒等它的那次调用
+
+        :param echo: 回执带回来的 echo
+        :param result: 整个 result 帧
+        """
+        if future := self._waiting.pop(echo, None):
+            if not future.done():
+                future.set_result(result)
+
+    def hold(self, task: asyncio.Task) -> None:
+        """把一个处理函数交给这个机器人管
+
+        :param task: 任务
+        """
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.discard)
+
+    def abort(self, reason: str) -> None:
+        """连接没了：挂着的调用立刻失败，还在跑的处理函数取消掉
+
+        处理函数拿着的是一条已经死了的连接，接着跑只会往空处撞。
+
+        :param reason: 失败原因
+        """
+        waiting, running = len(self._waiting), len(self._tasks)
+
+        for future in self._waiting.values():
+            if not future.done():
+                future.set_exception(CallError(0, {'error': reason}))
+        self._waiting.clear()
+
+        for task in list(self._tasks):
+            task.cancel()
+
+        if waiting or running:
+            logger.debug('{} 中止了 {} 次挂起调用、{} 个处理函数', self, waiting, running)