✨ AttachmentSegment 支持传递 url，添加获取 attachments 的方法 (#69)

* 🐛 修复 #43 转发图片消息时附件不可发送

* ♻️ 使用 yarl.URL 解析附件地址

Author: shoucandanghehe

nonebot/adapters/discord/bot.py

@@ -1,14 +1,18 @@
-from typing import TYPE_CHECKING, Any, Optional, Union
+from http import HTTPStatus
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 from typing_extensions import override
 
 from nonebot.adapters import Bot as BaseBot
 
+from nonebot.drivers import Request
 from nonebot.message import handle_event
+from yarl import URL
 
 from .api import (
     UNSET,
     AllowedMention,
     ApiClient,
+    File,
     InteractionCallbackMessage,
     InteractionCallbackType,
     InteractionResponse,
@@ -28,6 +32,10 @@
     from .adapter import Adapter
 
 
+DISCORD_ATTACHMENT_HOSTS = {"cdn.discordapp.com", "media.discordapp.net"}
+AttachmentFetchOnError = Literal["raise", "skip"]
+
+
 async def _check_reply(bot: "Bot", event: MessageEvent) -> None:
     message_reference = event.message_reference
     if message_reference is UNSET:
@@ -175,6 +183,138 @@ async def handle_event(self, event: Event) -> None:
             _check_at_me(self, event)
         await handle_event(self, event)
 
+    async def fetch_attachments(  # noqa: PLR0913
+        self,
+        message: Union[str, Message, MessageSegment],
+        *,
+        allowed_hosts: Optional[set[str]] = None,
+        require_https: bool = True,
+        timeout: Optional[float] = None,
+        max_bytes: Optional[int] = None,
+        prefer_proxy_url: bool = True,
+        on_error: AttachmentFetchOnError = "raise",
+    ) -> Message:
+        message = MessageSegment.text(message) if isinstance(message, str) else message
+        message = message if isinstance(message, Message) else Message(message)
+        new = message.clone()
+
+        if allowed_hosts is None:
+            allowed_hosts = DISCORD_ATTACHMENT_HOSTS
+
+        attachment_segments = new["attachment"] or []
+        for index, attachment in enumerate(attachment_segments):
+            if attachment.data["file"] is not None:
+                continue
+
+            url = self._pick_attachment_url(
+                attachment,
+                allowed_hosts=allowed_hosts,
+                require_https=require_https,
+                prefer_proxy_url=prefer_proxy_url,
+            )
+            if url is None:
+                if on_error == "raise":
+                    msg = (
+                        f"Attachment segment at index {index} has no fetchable "
+                        "url/proxy_url"
+                    )
+                    raise ValueError(msg)
+                continue
+
+            content = await self._fetch_attachment_content(
+                url,
+                timeout=timeout,
+                max_bytes=max_bytes,
+            )
+            if content is None:
+                if on_error == "raise":
+                    msg = (
+                        f"Failed to fetch attachment content for segment "
+                        f"at index {index} from URL {url}"
+                    )
+                    raise ValueError(msg)
+                continue
+
+            attachment.data["file"] = File(
+                filename=attachment.data["attachment"].filename,
+                content=content,
+            )
+
+        return new
+
+    @staticmethod
+    def _pick_attachment_url(
+        attachment: MessageSegment,
+        *,
+        allowed_hosts: set[str],
+        require_https: bool,
+        prefer_proxy_url: bool,
+    ) -> Optional[str]:
+        urls = []
+        if prefer_proxy_url:
+            urls.extend(
+                [
+                    attachment.data.get("proxy_url"),
+                    attachment.data.get("url"),
+                ]
+            )
+        else:
+            urls.extend(
+                [
+                    attachment.data.get("url"),
+                    attachment.data.get("proxy_url"),
+                ]
+            )
+
+        for candidate in urls:
+            if isinstance(candidate, str) and Bot._is_supported_attachment_url(
+                candidate,
+                allowed_hosts=allowed_hosts,
+                require_https=require_https,
+            ):
+                return candidate
+
+        return None
+
+    @staticmethod
+    def _is_supported_attachment_url(
+        url: str, *, allowed_hosts: set[str], require_https: bool
+    ) -> bool:
+        parsed = URL(url)
+        scheme_ok = parsed.scheme == "https" if require_https else bool(parsed.scheme)
+        return (
+            scheme_ok and isinstance(parsed.host, str) and parsed.host in allowed_hosts
+        )
+
+    async def _fetch_attachment_content(
+        self,
+        url: str,
+        *,
+        timeout: Optional[float],
+        max_bytes: Optional[int],
+    ) -> Optional[bytes]:
+        try:
+            request = Request(
+                method="GET",
+                url=url,
+                timeout=timeout or self.adapter.discord_config.discord_api_timeout,
+                proxy=self.adapter.discord_config.discord_proxy,
+            )
+            response = await self.adapter.request(request)
+            if response.status_code != HTTPStatus.OK or not response.content:
+                return None
+            content = (
+                response.content.encode()
+                if isinstance(response.content, str)
+                else response.content
+            )
+            if max_bytes is not None and len(content) > max_bytes:
+                return None
+            return content  # noqa: TRY300
+        except Exception as e:
+            log("DEBUG", f"Failed to fetch attachment content from URL {url}: {e!r}", e)
+            return None
+
     async def send_to(
         self,
         channel_id: SnowflakeType,
@@ -183,6 +323,9 @@ async def send_to(
         nonce: Union[int, str, None] = None,
         allowed_mentions: Optional[AllowedMention] = None,
     ) -> MessageGet:
+        message = MessageSegment.text(message) if isinstance(message, str) else message
+        message = message if isinstance(message, Message) else Message(message)
+        message = message.sendable()
         message_data = parse_message(message)
 
         return await self.create_message(
@@ -222,6 +365,8 @@ async def send(
             message model
         """
         message = MessageSegment.text(message) if isinstance(message, str) else message
+        message = message if isinstance(message, Message) else Message(message)
+        message = message.sendable()
         if isinstance(event, InteractionCreateEvent):
             message_data = parse_message(message)
             response = InteractionResponse(
@@ -250,7 +395,6 @@ async def send(
         if not isinstance(event, MessageEvent) or not event.channel_id or not event.id:
             msg = "Event cannot be replied to!"
             raise RuntimeError(msg)
-        message = message if isinstance(message, Message) else Message(message)
         if mention_sender or at_sender:
             message.insert(0, MessageSegment.mention_user(event.user_id))
         if reply_message:

nonebot/adapters/discord/message.py

@@ -1,4 +1,5 @@
 from collections.abc import Iterable
+from copy import deepcopy
 from dataclasses import dataclass
 import datetime
 import re
@@ -61,6 +62,9 @@ def attachment(
         file: Union[str, File, AttachmentSend],
         description: Optional[str] = None,
         content: Optional[bytes] = None,
+        *,
+        url: Optional[str] = None,
+        proxy_url: Optional[str] = None,
     ) -> "AttachmentSegment":
         if isinstance(file, str):
             _filename = file
@@ -85,6 +89,8 @@ def attachment(
                         filename=_filename, description=_description
                     ),
                     "file": None,
+                    "url": url,
+                    "proxy_url": proxy_url,
                 },
             )
         return AttachmentSegment(
@@ -98,6 +104,8 @@ def attachment(
                     if isinstance(_filename, str)
                     else None
                 ),
+                "url": url,
+                "proxy_url": proxy_url,
             },
         )
 
@@ -475,6 +483,8 @@ def _validate(cls, value) -> Self:  # noqa: ANN001
 class AttachmentData(TypedDict):
     attachment: AttachmentSend
     file: Optional[File]
+    url: Optional[str]
+    proxy_url: Optional[str]
 
 
 @dataclass
@@ -507,6 +517,17 @@ def _validate(cls, value) -> Self:  # noqa: ANN001
             file, File
         ):
             instance.data["file"] = type_validate_python(File, file)
+        url = instance.data.get("url")
+        if url is not None and not isinstance(url, str):
+            msg = f"Expected str for AttachmentSegment.data['url'], got {type(url)}"
+            raise TypeError(msg)
+        proxy_url = instance.data.get("proxy_url")
+        if proxy_url is not None and not isinstance(proxy_url, str):
+            msg = (
+                "Expected str for AttachmentSegment.data['proxy_url'], "
+                f"got {type(proxy_url)}"
+            )
+            raise TypeError(msg)
         return instance
 
 
@@ -676,7 +697,9 @@ def from_guild_message(cls, message: MessageGet) -> "Message":
                             if isinstance(attachment.description, str)
                             else None
                         ),
-                    )
+                    ),
+                    url=attachment.url,
+                    proxy_url=attachment.proxy_url,
                 )
                 for attachment in message.attachments
             )
@@ -706,6 +729,29 @@ def extract_content(self) -> str:
             )
         )
 
+    def clone(self) -> "Message":
+        new = self.__class__()
+        for segment in self:
+            new.append(
+                type_validate_python(
+                    MessageSegment,
+                    {
+                        "type": segment.type,
+                        "data": deepcopy(segment.data),
+                    },
+                )
+            )
+        return new
+
+    def sendable(self) -> "Message":
+        new = self.clone()
+        attachments_segment = new["attachment"] or None
+        if attachments_segment is not None:
+            for index, attachment in enumerate(attachments_segment):
+                if attachment.data["file"] is None:
+                    raise ValueError(_get_unsendable_attachment_msg(index, attachment))
+        return new
+
 
 def parse_message(message: Union[Message, MessageSegment, str]) -> dict[str, Any]:
     message = MessageSegment.text(message) if isinstance(message, str) else message
@@ -733,17 +779,8 @@ def parse_message(message: Union[Message, MessageSegment, str]) -> dict[str, Any
                 layout_type=poll.layout_type,
             )
 
-    attachments = None
-    files = None
-    if attachments_segment := (message["attachment"] or None):
-        attachments = [
-            attachment.data["attachment"] for attachment in attachments_segment
-        ]
-        files = [
-            attachment.data["file"]
-            for attachment in attachments_segment
-            if attachment.data["file"] is not None
-        ]
+    attachments, files = extract_attachments(message)
+
     return {
         k: v
         for k, v in {
@@ -758,3 +795,38 @@ def parse_message(message: Union[Message, MessageSegment, str]) -> dict[str, Any
         }.items()
         if v is not None
     }
+
+
+def extract_attachments(
+    message: Message,
+) -> tuple[Optional[list[AttachmentSend]], Optional[list[File]]]:
+    attachments_segment = message["attachment"] or None
+    if not attachments_segment:
+        return None, None
+
+    attachments_list: list[AttachmentSend] = []
+    files_list: list[File] = []
+    for index, attachment in enumerate(attachments_segment):
+        file = attachment.data["file"]
+        if file is None:
+            raise ValueError(_get_unsendable_attachment_msg(index, attachment))
+        attachments_list.append(attachment.data["attachment"])
+        files_list.append(file)
+
+    attachments = attachments_list or None
+    files = files_list or None
+    return attachments, files
+
+
+def _get_unsendable_attachment_msg(index: int, attachment: MessageSegment) -> str:
+    if attachment.data.get("url") or attachment.data.get("proxy_url"):
+        return (
+            f"Attachment segment at index {index} is not sendable because file "
+            "content is missing; call "
+            "`await bot.fetch_attachments(message)` first"
+        )
+    return (
+        f"Attachment segment at index {index} is not sendable because file "
+        "content is missing; provide `content=` in "
+        "MessageSegment.attachment(...)"
+    )