From aa2529c0273e41f6753309532dfcd1ee38808f6c Mon Sep 17 00:00:00 2001 From: dolfies Date: Mon, 24 Apr 2023 20:35:13 -0400 Subject: [PATCH] Implement GCP cloud uploads --- discord/abc.py | 58 +++++++++-- discord/ext/commands/context.py | 10 +- discord/file.py | 167 ++++++++++++++++++++++++++------ discord/http.py | 31 ++++-- discord/message.py | 29 +++--- discord/types/message.py | 16 +++ docs/api.rst | 7 ++ 7 files changed, 250 insertions(+), 68 deletions(-) diff --git a/discord/abc.py b/discord/abc.py index 333594eba..bfca8a5f8 100644 --- a/discord/abc.py +++ b/discord/abc.py @@ -54,7 +54,7 @@ from .mentions import AllowedMentions from .permissions import PermissionOverwrite, Permissions from .role import Role from .invite import Invite -from .file import File +from .file import File, CloudFile from .http import handle_message_parameters from .voice_client import VoiceClient, VoiceProtocol from .sticker import GuildSticker, StickerItem @@ -79,6 +79,7 @@ if TYPE_CHECKING: from .client import Client from .user import ClientUser, User from .asset import Asset + from .file import _FileBase from .state import ConnectionState from .guild import Guild from .member import Member @@ -1570,13 +1571,52 @@ class Messageable: async def _get_channel(self) -> MessageableChannel: raise NotImplementedError + async def upload_files(self, *files: File) -> List[CloudFile]: + r"""|coro| + + Pre-uploads files to Discord's GCP bucket for use with :meth:`send`. + + This method is useful if you have local files that you want to upload and + reuse multiple times. + + Parameters + ------------ + \*files: :class:`~discord.File` + A list of files to upload. Must be a maximum of 10. + + Raises + ------- + ~discord.HTTPException + Uploading the files failed. + ~discord.Forbidden + You do not have the proper permissions to upload files. + + Returns + -------- + List[:class:`~discord.CloudFile`] + The files that were uploaded. These can be used in lieu + of normal :class:`~discord.File`\s in :meth:`send`. + """ + if not files: + return [] + + state = self._state + channel = await self._get_channel() + + mapped_files = {i: f for i, f in enumerate(files)} + data = await self._state.http.get_attachment_urls(channel.id, [f.to_upload_dict(i) for i, f in mapped_files.items()]) + return [ + await CloudFile.from_file(state=state, data=uploaded, file=mapped_files[int(uploaded.get('id', 11))]) + for uploaded in data['attachments'] + ] + @overload async def send( self, content: Optional[str] = ..., *, tts: bool = ..., - file: File = ..., + file: _FileBase = ..., stickers: Sequence[Union[GuildSticker, StickerItem]] = ..., delete_after: float = ..., nonce: Union[str, int] = ..., @@ -1594,7 +1634,7 @@ class Messageable: content: Optional[str] = ..., *, tts: bool = ..., - files: Sequence[File] = ..., + files: Sequence[_FileBase] = ..., stickers: Sequence[Union[GuildSticker, StickerItem]] = ..., delete_after: float = ..., nonce: Union[str, int] = ..., @@ -1612,7 +1652,7 @@ class Messageable: content: Optional[str] = ..., *, tts: bool = ..., - file: File = ..., + file: _FileBase = ..., stickers: Sequence[Union[GuildSticker, StickerItem]] = ..., delete_after: float = ..., nonce: Union[str, int] = ..., @@ -1630,7 +1670,7 @@ class Messageable: content: Optional[str] = ..., *, tts: bool = ..., - files: Sequence[File] = ..., + files: Sequence[_FileBase] = ..., stickers: Sequence[Union[GuildSticker, StickerItem]] = ..., delete_after: float = ..., nonce: Union[str, int] = ..., @@ -1647,8 +1687,8 @@ class Messageable: content: Optional[str] = None, *, tts: bool = False, - file: Optional[File] = None, - files: Optional[Sequence[File]] = None, + file: Optional[_FileBase] = None, + files: Optional[Sequence[_FileBase]] = None, stickers: Optional[Sequence[Union[GuildSticker, StickerItem]]] = None, delete_after: Optional[float] = None, nonce: Optional[Union[str, int]] = MISSING, @@ -1680,9 +1720,9 @@ class Messageable: The content of the message to send. tts: :class:`bool` Indicates if the message should be sent using text-to-speech. - file: :class:`~discord.File` + file: Union[:class:`~discord.File`, :class:`~discord.CloudFile`] The file to upload. - files: List[:class:`~discord.File`] + files: List[Union[:class:`~discord.File`, :class:`~discord.CloudFile`]] A list of files to upload. Must be a maximum of 10. nonce: :class:`int` The nonce to use for sending this message. If the message was successfully sent, diff --git a/discord/ext/commands/context.py b/discord/ext/commands/context.py index fd8076da1..145a4f828 100644 --- a/discord/ext/commands/context.py +++ b/discord/ext/commands/context.py @@ -51,7 +51,7 @@ if TYPE_CHECKING: from discord.abc import MessageableChannel from discord.commands import MessageCommand - from discord.file import File + from discord.file import _FileBase from discord.guild import Guild from discord.member import Member from discord.mentions import AllowedMentions @@ -433,7 +433,7 @@ class Context(discord.abc.Messageable, Generic[BotT]): content: Optional[str] = ..., *, tts: bool = ..., - file: File = ..., + file: _FileBase = ..., stickers: Sequence[Union[GuildSticker, StickerItem]] = ..., delete_after: float = ..., nonce: Union[str, int] = ..., @@ -451,7 +451,7 @@ class Context(discord.abc.Messageable, Generic[BotT]): content: Optional[str] = ..., *, tts: bool = ..., - files: Sequence[File] = ..., + files: Sequence[_FileBase] = ..., stickers: Sequence[Union[GuildSticker, StickerItem]] = ..., delete_after: float = ..., nonce: Union[str, int] = ..., @@ -469,7 +469,7 @@ class Context(discord.abc.Messageable, Generic[BotT]): content: Optional[str] = ..., *, tts: bool = ..., - file: File = ..., + file: _FileBase = ..., stickers: Sequence[Union[GuildSticker, StickerItem]] = ..., delete_after: float = ..., nonce: Union[str, int] = ..., @@ -487,7 +487,7 @@ class Context(discord.abc.Messageable, Generic[BotT]): content: Optional[str] = ..., *, tts: bool = ..., - files: Sequence[File] = ..., + files: Sequence[_FileBase] = ..., stickers: Sequence[Union[GuildSticker, StickerItem]] = ..., delete_after: float = ..., nonce: Union[str, int] = ..., diff --git a/discord/file.py b/discord/file.py index abb8e707b..965116470 100644 --- a/discord/file.py +++ b/discord/file.py @@ -28,15 +28,21 @@ from base64 import b64encode from hashlib import md5 import io import os -from typing import Any, Dict, Optional, Tuple, Union +import yarl +from typing import Any, Dict, Optional, Tuple, Union, TYPE_CHECKING from .utils import MISSING, cached_slot_property -# fmt: off +if TYPE_CHECKING: + from typing_extensions import Self + + from .state import ConnectionState + from .types.message import CloudAttachment as CloudAttachmentPayload, UploadedAttachment as UploadedAttachmentPayload + __all__ = ( 'File', + 'CloudFile', ) -# fmt: on def _strip_spoiler(filename: str) -> Tuple[str, bool]: @@ -47,7 +53,48 @@ def _strip_spoiler(filename: str) -> Tuple[str, bool]: return stripped, spoiler -class File: +class _FileBase: + __slots__ = ('_filename', 'spoiler', 'description') + + def __init__(self, filename: str, *, spoiler: bool = False, description: Optional[str] = None): + self._filename, filename_spoiler = _strip_spoiler(filename) + if spoiler is MISSING: + spoiler = filename_spoiler + + self.spoiler: bool = spoiler + self.description: Optional[str] = description + + @property + def filename(self) -> str: + """:class:`str`: The filename to display when uploading to Discord. + If this is not given then it defaults to ``fp.name`` or if ``fp`` is + a string then the ``filename`` will default to the string given. + """ + return 'SPOILER_' + self._filename if self.spoiler else self._filename + + @filename.setter + def filename(self, value: str) -> None: + self._filename, self.spoiler = _strip_spoiler(value) + + def to_dict(self, index: int) -> Dict[str, Any]: + payload = { + 'id': str(index), + 'filename': self.filename, + } + + if self.description is not None: + payload['description'] = self.description + + return payload + + def reset(self, *, seek: Union[int, bool] = True) -> None: + return + + def close(self) -> None: + return + + +class File(_FileBase): r"""A parameter object used for :meth:`abc.Messageable.send` for sending file objects. @@ -79,7 +126,7 @@ class File: .. versionadded:: 2.0 """ - __slots__ = ('fp', '_filename', 'spoiler', 'description', '_original_pos', '_owner', '_closer', '_cs_md5') + __slots__ = ('fp', '_original_pos', '_owner', '_closer', '_cs_md5', '_cs_size') def __init__( self, @@ -112,24 +159,7 @@ class File: else: filename = getattr(fp, 'name', 'untitled') - self._filename, filename_spoiler = _strip_spoiler(filename) - if spoiler is MISSING: - spoiler = filename_spoiler - - self.spoiler: bool = spoiler - self.description: Optional[str] = description - - @property - def filename(self) -> str: - """:class:`str`: The filename to display when uploading to Discord. - If this is not given then it defaults to ``fp.name`` or if ``fp`` is - a string then the ``filename`` will default to the string given. - """ - return 'SPOILER_' + self._filename if self.spoiler else self._filename - - @filename.setter - def filename(self, value: str) -> None: - self._filename, self.spoiler = _strip_spoiler(value) + super().__init__(filename, spoiler=spoiler, description=description) @cached_slot_property('_cs_md5') def md5(self) -> str: @@ -138,6 +168,17 @@ class File: finally: self.reset() + @cached_slot_property('_cs_size') + def size(self) -> int: + return os.fstat(self.fp.fileno()).st_size + + def to_upload_dict(self, index: int) -> UploadedAttachmentPayload: + return { + 'id': str(index), + 'filename': self.filename, + 'file_size': self.size, + } + def reset(self, *, seek: Union[int, bool] = True) -> None: # The `seek` parameter is needed because # the retry-loop is iterated over multiple times @@ -155,13 +196,79 @@ class File: if self._owner: self._closer() - def to_dict(self, index: int) -> Dict[str, Any]: - payload = { - 'id': index, - 'filename': self.filename, - } - if self.description is not None: - payload['description'] = self.description +class CloudFile(_FileBase): + """A parameter object used for :meth:`abc.Messageable.send` + for sending file objects that have been pre-uploaded to Discord's GCP bucket. + + .. note:: + + Unlike :class:`File`, this class is not directly user-constructable, however + it can be reused multiple times in :meth:`abc.Messageable.send`. + + To construct it, see :meth:`abc.Messageable.upload_files`. + + .. versionadded:: 2.1 + + Attributes + ----------- + url: :class:`str` + The upload URL of the file. + .. note:: + + This URL cannot be used to download the file, + it is merely used to send the file to Discord. + upload_filename: :class:`str` + The filename that Discord has assigned to the file. + spoiler: :class:`bool` + Whether the attachment is a spoiler. If left unspecified, the :attr:`~CloudFile.filename` is used + to determine if the file is a spoiler. + description: Optional[:class:`str`] + The file description to display, currently only supported for images. + """ + + __slots__ = ('url', 'upload_filename', '_state') + + def __init__( + self, + url: str, + filename: str, + upload_filename: str, + *, + spoiler: bool = MISSING, + description: Optional[str] = None, + state: ConnectionState, + ): + super().__init__(filename, spoiler=spoiler, description=description) + self.url = url + self.upload_filename = upload_filename + self._state = state + + @classmethod + async def from_file(cls, *, file: File, state: ConnectionState, data: CloudAttachmentPayload) -> Self: + await state.http.upload_to_cloud(data['upload_url'], file) + return cls(data['upload_url'], file._filename, data['upload_filename'], description=file.description, state=state) + + @property + def upload_id(self) -> str: + """:class:`str`: The upload ID of the file.""" + url = yarl.URL(self.url) + return url.query['upload_id'] + + def to_dict(self, index: int) -> Dict[str, Any]: + payload = super().to_dict(index) + payload['uploaded_filename'] = self.upload_filename return payload + + async def delete(self) -> None: + """|coro| + + Deletes the uploaded file from Discord's GCP bucket. + + Raises + ------- + HTTPException + Deleting the file failed. + """ + await self._state.http.delete_attachment(self.upload_filename) diff --git a/discord/http.py b/discord/http.py index 4a7d8233a..d733e991d 100644 --- a/discord/http.py +++ b/discord/http.py @@ -63,7 +63,7 @@ from .errors import ( GatewayNotFound, CaptchaRequired, ) -from .file import File +from .file import _FileBase, File from .tracking import ContextProperties from . import utils from .mentions import AllowedMentions @@ -75,7 +75,6 @@ if TYPE_CHECKING: from .channel import TextChannel, DMChannel, GroupChannel, PartialMessageable, VoiceChannel, ForumChannel from .handlers import CaptchaHandler from .threads import Thread - from .file import File from .mentions import AllowedMentions from .message import Attachment, Message from .flags import MessageFlags @@ -166,11 +165,11 @@ def handle_message_parameters( tts: bool = False, nonce: Optional[Union[int, str]] = MISSING, flags: MessageFlags = MISSING, - file: File = MISSING, - files: Sequence[File] = MISSING, + file: _FileBase = MISSING, + files: Sequence[_FileBase] = MISSING, embed: Optional[Embed] = MISSING, embeds: Sequence[Embed] = MISSING, - attachments: Sequence[Union[Attachment, File]] = MISSING, + attachments: Sequence[Union[Attachment, _FileBase]] = MISSING, allowed_mentions: Optional[AllowedMentions] = MISSING, message_reference: Optional[message.MessageReference] = MISSING, stickers: Optional[SnowflakeList] = MISSING, @@ -249,13 +248,13 @@ def handle_message_parameters( if attachments is MISSING: attachments = files else: - files = [a for a in attachments if isinstance(a, File)] + files = [a for a in attachments if isinstance(a, _FileBase)] if attachments is not MISSING: file_index = 0 attachments_payload = [] for attachment in attachments: - if isinstance(attachment, File): + if isinstance(attachment, _FileBase): attachments_payload.append(attachment.to_dict(file_index)) file_index += 1 else: @@ -269,11 +268,13 @@ def handle_message_parameters( } payload.update(channel_payload) + # Legacy uploading multipart = [] - if files: + to_upload = [file for file in files if isinstance(file, File)] + if to_upload: multipart.append({'name': 'payload_json', 'value': utils._to_json(payload)}) payload = None - for index, file in enumerate(files): + for index, file in enumerate(to_upload): multipart.append( { 'name': f'files[{index}]', @@ -283,7 +284,7 @@ def handle_message_parameters( } ) - return MultipartParameters(payload=payload, multipart=multipart, files=files) + return MultipartParameters(payload=payload, multipart=multipart, files=to_upload) def _gen_accept_encoding_header(): @@ -1326,6 +1327,16 @@ class HTTPClient: def ack_pins(self, channel_id: Snowflake) -> Response[None]: return self.request(Route('POST', '/channels/{channel_id}/pins/ack', channel_id=channel_id)) + def get_attachment_urls( + self, channel_id: Snowflake, attachments: List[message.UploadedAttachment] + ) -> Response[message.CloudAttachments]: + payload = {'files': attachments} + + return self.request(Route('POST', '/channels/{channel_id}/attachments', channel_id=channel_id), json=payload) + + def delete_attachment(self, uploaded_filename: str) -> Response[None]: + return self.request(Route('DELETE', '/attachments/{uploaded_filename}', uploaded_filename=uploaded_filename)) + # Member management def kick(self, user_id: Snowflake, guild_id: Snowflake, reason: Optional[str] = None) -> Response[None]: diff --git a/discord/message.py b/discord/message.py index 21149f9d4..57a0490d1 100644 --- a/discord/message.py +++ b/discord/message.py @@ -99,6 +99,7 @@ if TYPE_CHECKING: from .abc import Snowflake from .abc import GuildChannel, MessageableChannel from .components import ActionRow, ActionRowChildComponentType + from .file import _FileBase from .state import ConnectionState from .mentions import AllowedMentions from .sticker import GuildSticker @@ -745,7 +746,7 @@ class PartialMessage(Hashable): self, *, content: Optional[str] = ..., - attachments: Sequence[Union[Attachment, File]] = ..., + attachments: Sequence[Union[Attachment, _FileBase]] = ..., delete_after: Optional[float] = ..., allowed_mentions: Optional[AllowedMentions] = ..., ) -> Message: @@ -756,7 +757,7 @@ class PartialMessage(Hashable): self, *, content: Optional[str] = ..., - attachments: Sequence[Union[Attachment, File]] = ..., + attachments: Sequence[Union[Attachment, _FileBase]] = ..., delete_after: Optional[float] = ..., allowed_mentions: Optional[AllowedMentions] = ..., ) -> Message: @@ -765,7 +766,7 @@ class PartialMessage(Hashable): async def edit( self, content: Optional[str] = MISSING, - attachments: Sequence[Union[Attachment, File]] = MISSING, + attachments: Sequence[Union[Attachment, _FileBase]] = MISSING, delete_after: Optional[float] = None, allowed_mentions: Optional[AllowedMentions] = MISSING, ) -> Message: @@ -787,7 +788,7 @@ class PartialMessage(Hashable): content: Optional[:class:`str`] The new content to replace the message with. Could be ``None`` to remove the content. - attachments: List[Union[:class:`Attachment`, :class:`File`]] + attachments: List[Union[:class:`Attachment`, :class:`File`, :class:`CloudFile`]] A list of attachments to keep in the message as well as new files to upload. If ``[]`` is passed then all attachments are removed. @@ -1171,7 +1172,7 @@ class PartialMessage(Hashable): content: Optional[str] = ..., *, tts: bool = ..., - file: File = ..., + file: _FileBase = ..., stickers: Sequence[Union[GuildSticker, StickerItem]] = ..., delete_after: float = ..., nonce: Union[str, int] = ..., @@ -1188,7 +1189,7 @@ class PartialMessage(Hashable): content: Optional[str] = ..., *, tts: bool = ..., - files: Sequence[File] = ..., + files: Sequence[_FileBase] = ..., stickers: Sequence[Union[GuildSticker, StickerItem]] = ..., delete_after: float = ..., nonce: Union[str, int] = ..., @@ -1205,7 +1206,7 @@ class PartialMessage(Hashable): content: Optional[str] = ..., *, tts: bool = ..., - file: File = ..., + file: _FileBase = ..., stickers: Sequence[Union[GuildSticker, StickerItem]] = ..., delete_after: float = ..., nonce: Union[str, int] = ..., @@ -1222,7 +1223,7 @@ class PartialMessage(Hashable): content: Optional[str] = ..., *, tts: bool = ..., - files: Sequence[File] = ..., + files: Sequence[_FileBase] = ..., stickers: Sequence[Union[GuildSticker, StickerItem]] = ..., delete_after: float = ..., nonce: Union[str, int] = ..., @@ -2122,7 +2123,7 @@ class Message(PartialMessage, Hashable): self, *, content: Optional[str] = ..., - attachments: Sequence[Union[Attachment, File]] = ..., + attachments: Sequence[Union[Attachment, _FileBase]] = ..., suppress: bool = ..., delete_after: Optional[float] = ..., allowed_mentions: Optional[AllowedMentions] = ..., @@ -2134,7 +2135,7 @@ class Message(PartialMessage, Hashable): self, *, content: Optional[str] = ..., - attachments: Sequence[Union[Attachment, File]] = ..., + attachments: Sequence[Union[Attachment, _FileBase]] = ..., suppress: bool = ..., delete_after: Optional[float] = ..., allowed_mentions: Optional[AllowedMentions] = ..., @@ -2144,7 +2145,7 @@ class Message(PartialMessage, Hashable): async def edit( self, content: Optional[str] = MISSING, - attachments: Sequence[Union[Attachment, File]] = MISSING, + attachments: Sequence[Union[Attachment, _FileBase]] = MISSING, suppress: bool = False, delete_after: Optional[float] = None, allowed_mentions: Optional[AllowedMentions] = MISSING, @@ -2170,7 +2171,7 @@ class Message(PartialMessage, Hashable): content: Optional[:class:`str`] The new content to replace the message with. Could be ``None`` to remove the content. - attachments: List[Union[:class:`Attachment`, :class:`File`]] + attachments: List[Union[:class:`Attachment`, :class:`File`, :class:`CloudFile`]] A list of attachments to keep in the message as well as new files to upload. If ``[]`` is passed then all attachments are removed. @@ -2238,7 +2239,7 @@ class Message(PartialMessage, Hashable): return message - async def add_files(self, *files: File) -> Message: + async def add_files(self, *files: _FileBase) -> Message: r"""|coro| Adds new files to the end of the message attachments. @@ -2247,7 +2248,7 @@ class Message(PartialMessage, Hashable): Parameters ----------- - \*files: :class:`File` + \*files: Union[:class:`File`, :class:`CloudFile`] New files to add to the message. Raises diff --git a/discord/types/message.py b/discord/types/message.py index e5c7ca315..bbed7c624 100644 --- a/discord/types/message.py +++ b/discord/types/message.py @@ -185,3 +185,19 @@ MessageSearchHasType = Literal[ ] MessageSearchSortType = Literal['timestamp', 'relevance'] MessageSearchSortOrder = Literal['desc', 'asc'] + + +class UploadedAttachment(TypedDict): + id: NotRequired[Snowflake] + filename: str + file_size: int + + +class CloudAttachment(TypedDict): + id: NotRequired[Snowflake] + upload_url: str + upload_filename: str + + +class CloudAttachments(TypedDict): + attachments: List[CloudAttachment] diff --git a/docs/api.rst b/docs/api.rst index 177af1c68..b20b69f7f 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -7637,6 +7637,13 @@ File .. autoclass:: File() :members: + :inherited-members: + +.. attributetable:: CloudFile + +.. autoclass:: CloudFile() + :members: + :inherited-members: Colour ~~~~~~