onebot_v12_segment.py

from io import BytesIO
from pathlib import Path
from typing import Union
from base64 import b64decode
from urllib.parse import urlparse

from nonebot.log import logger
from nonebot.adapters.onebot.v12 import Bot, MessageSegment


class FileV12:
    """OneBot V12 含 ``file_id`` 消息段生成"""

    def __init__(
        self,
        bot: Bot,
        file: Union[str, bytes, BytesIO, Path],
        name: str = "",
    ) -> None:
        self.bot = bot  # TODO: use nonebot.get_bot
        self.file = file
        self.name = name

    @property
    async def file_id(self) -> str:
        return await self.get_file_id(self.bot, self.file, self.name)

    @staticmethod
    async def get_file_id(
        bot: Bot, file: Union[str, bytes, BytesIO, Path], name: str = ""
    ) -> str:
        """通过 OneBot V12 上传文件接口获取 ``file_id``"""

        if isinstance(file, BytesIO):
            file = file.getvalue()
        if isinstance(file, str) and file.startswith("file:///"):
            file = Path(file.lstrip("file:///"))
        if isinstance(file, str) and file.startswith("base64://"):
            file = b64decode(file.lstrip("base64://"))

        if isinstance(file, bytes):
            params = {"type": "data", "name": name, "data": file}
        elif isinstance(file, Path):
            params = {
                "type": "path",
                "name": file.name,
                "path": file.resolve().as_uri(),
            }
        elif isinstance(file, str) and file.startswith(("http://", "https://")):
            params = {
                "type": "url",
                "name": urlparse(file).path.split("/")[-1],
                "url": file,
            }
        else:
            logger.warning("Unsupported type of file!")
            return ""

        upload_res = await bot.upload_file(**params)
        if upload_res and upload_res.get("file_id"):
            return upload_res["file_id"]
        logger.warning("File upload failed!")
        return ""

    async def as_image(self) -> MessageSegment:
        """生成 OneBot V12 图片消息段 ``MessageSegment.image()``"""
        return MessageSegment("image", {"file_id": await self.file_id})

    async def as_voice(self) -> MessageSegment:
        """生成 OneBot V12 语音消息段 ``MessageSegment.voice()``"""
        return MessageSegment("voice", {"file_id": await self.file_id})

    async def as_audio(self) -> MessageSegment:
        """生成 OneBot V12 音频消息段 ``MessageSegment.audio()``"""
        return MessageSegment("audio", {"file_id": await self.file_id})

    async def as_video(self) -> MessageSegment:
        """生成 OneBot V12 视频消息段 ``MessageSegment.video()``"""
        return MessageSegment("video", {"file_id": await self.file_id})

    async def as_file(self) -> MessageSegment:
        """生成 OneBot V12 文件消息段 ``MessageSegment.file()``"""
        return MessageSegment("file", {"file_id": await self.file_id})


# Get MessageSegment.image() :
# await FileV12(bot, file).as_image()