From fed259a83ba69f8651f88c37a458b287737c44ea Mon Sep 17 00:00:00 2001 From: Rapptz Date: Sat, 17 Apr 2021 00:55:44 -0400 Subject: [PATCH] Refactor save() and read() into AssetMixin --- discord/asset.py | 149 +++++++++++++++++++-------------------- discord/emoji.py | 61 +--------------- discord/message.py | 2 +- discord/partial_emoji.py | 72 ++----------------- 4 files changed, 78 insertions(+), 206 deletions(-) diff --git a/discord/asset.py b/discord/asset.py index bf7cd6ca1..3785dca2f 100644 --- a/discord/asset.py +++ b/discord/asset.py @@ -26,7 +26,7 @@ from __future__ import annotations import io import os -from typing import BinaryIO, Literal, TYPE_CHECKING, Tuple, Union +from typing import Any, Literal, Optional, TYPE_CHECKING, Tuple, Union from .errors import DiscordException from .errors import InvalidArgument from . import utils @@ -45,7 +45,76 @@ VALID_STATIC_FORMATS = frozenset({"jpeg", "jpg", "webp", "png"}) VALID_ASSET_FORMATS = VALID_STATIC_FORMATS | {"gif"} -class Asset: +class AssetMixin: + url: str + _state: Optional[Any] + + async def read(self) -> bytes: + """|coro| + + Retrieves the content of this asset as a :class:`bytes` object. + + Raises + ------ + DiscordException + There was no internal connection state. + HTTPException + Downloading the asset failed. + NotFound + The asset was deleted. + + Returns + ------- + :class:`bytes` + The content of the asset. + """ + if self._state is None: + raise DiscordException('Invalid state (no ConnectionState provided)') + + return await self._state.http.get_from_cdn(self.url) + + async def save(self, fp: Union[str, bytes, os.PathLike, io.BufferedIOBase], *, seek_begin: bool = True) -> int: + """|coro| + + Saves this asset into a file-like object. + + Parameters + ---------- + fp: Union[:class:`io.BufferedIOBase`, :class:`os.PathLike`] + The file-like object to save this attachment to or the filename + to use. If a filename is passed then a file is created with that + filename and used instead. + seek_begin: :class:`bool` + Whether to seek to the beginning of the file after saving is + successfully done. + + Raises + ------ + DiscordException + There was no internal connection state. + HTTPException + Downloading the asset failed. + NotFound + The asset was deleted. + + Returns + -------- + :class:`int` + The number of bytes written. + """ + + data = await self.read() + if isinstance(fp, io.BufferedIOBase): + written = fp.write(data) + if seek_begin: + fp.seek(0) + return written + else: + with open(fp, 'wb') as f: + return f.write(data) + + +class Asset(AssetMixin): """Represents a CDN asset on Discord. .. container:: operations @@ -153,19 +222,6 @@ class Asset: animated=False, ) - @classmethod - def _from_emoji(cls, state, emoji, *, format=None, static_format='png'): - if format is not None and format not in VALID_AVATAR_FORMATS: - raise InvalidArgument(f"format must be None or one of {VALID_AVATAR_FORMATS}") - if format == "gif" and not emoji.animated: - raise InvalidArgument("non animated emoji's do not support gif format") - if static_format not in VALID_STATIC_FORMATS: - raise InvalidArgument(f"static_format must be one of {VALID_STATIC_FORMATS}") - if format is None: - format = 'gif' if emoji.animated else static_format - - return cls(state, f'/emojis/{emoji.id}.{format}') - def __str__(self) -> str: return self._url @@ -332,66 +388,3 @@ class Asset: if self._animated: return self return self.with_format(format) - - async def read(self) -> bytes: - """|coro| - - Retrieves the content of this asset as a :class:`bytes` object. - - .. versionadded:: 1.1 - - Raises - ------ - DiscordException - There was no internal connection state. - HTTPException - Downloading the asset failed. - NotFound - The asset was deleted. - - Returns - ------- - :class:`bytes` - The content of the asset. - """ - if self._state is None: - raise DiscordException('Invalid state (no ConnectionState provided)') - - return await self._state.http.get_from_cdn(self.BASE + self._url) - - async def save(self, fp: Union[str, bytes, os.PathLike, BinaryIO], *, seek_begin: bool = True) -> int: - """|coro| - - Saves this asset into a file-like object. - - Parameters - ---------- - fp: Union[BinaryIO, :class:`os.PathLike`] - Same as in :meth:`Attachment.save`. - seek_begin: :class:`bool` - Same as in :meth:`Attachment.save`. - - Raises - ------ - DiscordException - There was no internal connection state. - HTTPException - Downloading the asset failed. - NotFound - The asset was deleted. - - Returns - -------- - :class:`int` - The number of bytes written. - """ - - data = await self.read() - if isinstance(fp, io.IOBase) and fp.writable(): - written = fp.write(data) - if seek_begin: - fp.seek(0) - return written - else: - with open(fp, 'wb') as f: - return f.write(data) diff --git a/discord/emoji.py b/discord/emoji.py index aba6d765e..adc646960 100644 --- a/discord/emoji.py +++ b/discord/emoji.py @@ -24,7 +24,7 @@ DEALINGS IN THE SOFTWARE. import io -from .asset import Asset +from .asset import Asset, AssetMixin from . import utils from .partial_emoji import _EmojiTag from .user import User @@ -33,7 +33,7 @@ __all__ = ( 'Emoji', ) -class Emoji(_EmojiTag): +class Emoji(_EmojiTag, AssetMixin): """Represents a custom emoji. Depending on the way this object was created, some of the attributes can @@ -221,60 +221,3 @@ class Emoji(_EmojiTag): roles = [role.id for role in roles] await self._state.http.edit_custom_emoji(self.guild.id, self.id, name=name, roles=roles, reason=reason) - async def read(self): - """|coro| - - Retrieves the content of this emoji as a :class:`bytes` object. - - .. versionadded:: 2.0 - - Raises - ------ - HTTPException - Downloading the emoji failed. - NotFound - The emoji was deleted. - - Returns - ------- - :class:`bytes` - The content of the emoji. - """ - return await self._state.http.get_from_cdn(self.url) - - async def save(self, fp, *, seek_begin=True): - """|coro| - - Saves this emoji into a file-like object. - - .. versionadded:: 2.0 - - Parameters - ---------- - fp: Union[BinaryIO, :class:`os.PathLike`] - Same as in :meth:`Attachment.save`. - seek_begin: :class:`bool` - Same as in :meth:`Attachment.save`. - - Raises - ------ - HTTPException - Downloading the emoji failed. - NotFound - The emoji was deleted. - - Returns - -------- - :class:`int` - The number of bytes written. - """ - - data = await self.read() - if isinstance(fp, io.IOBase) and fp.writable(): - written = fp.write(data) - if seek_begin: - fp.seek(0) - return written - else: - with open(fp, 'wb') as f: - return f.write(data) diff --git a/discord/message.py b/discord/message.py index e434562c2..df0b8243f 100644 --- a/discord/message.py +++ b/discord/message.py @@ -172,7 +172,7 @@ class Attachment(Hashable): The number of bytes written. """ data = await self.read(use_cached=use_cached) - if isinstance(fp, io.IOBase) and fp.writable(): + if isinstance(fp, io.BufferedIOBase): written = fp.write(data) if seek_begin: fp.seek(0) diff --git a/discord/partial_emoji.py b/discord/partial_emoji.py index 0c65dc341..82203bf53 100644 --- a/discord/partial_emoji.py +++ b/discord/partial_emoji.py @@ -24,7 +24,7 @@ DEALINGS IN THE SOFTWARE. import io -from .asset import Asset +from .asset import Asset, AssetMixin from .errors import DiscordException, InvalidArgument from . import utils @@ -35,7 +35,7 @@ __all__ = ( class _EmojiTag: __slots__ = () -class PartialEmoji(_EmojiTag): +class PartialEmoji(_EmojiTag, AssetMixin): """Represents a "partial" emoji. This model will be given in two scenarios: @@ -163,72 +163,8 @@ class PartialEmoji(_EmojiTag): fmt = 'gif' if self.animated else 'png' return f'{Asset.BASE}/emojis/{self.id}.{fmt}' - async def read(self): - """|coro| - - Retrieves the content of this emoji as a :class:`bytes` object. - - .. versionadded:: 2.0 - - Raises - ------ - DiscordException - There was no internal connection state. - InvalidArgument - The emoji isn't custom. - HTTPException - Downloading the emoji failed. - NotFound - The emoji was deleted. - - Returns - ------- - :class:`bytes` - The content of the emoji. - """ - if self._state is None: - raise DiscordException('Invalid state (no ConnectionState provided)') - + async def read(self) -> bytes: if self.is_unicode_emoji(): raise InvalidArgument('PartialEmoji is not a custom emoji') - return await self._state.http.get_from_cdn(self.url) - - async def save(self, fp, *, seek_begin=True): - """|coro| - - Saves this emoji into a file-like object. - - .. versionadded:: 2.0 - - Parameters - ---------- - fp: Union[BinaryIO, :class:`os.PathLike`] - Same as in :meth:`Attachment.save`. - seek_begin: :class:`bool` - Same as in :meth:`Attachment.save`. - - Raises - ------ - DiscordException - There was no internal connection state. - HTTPException - Downloading the emoji failed. - NotFound - The emoji was deleted. - - Returns - -------- - :class:`int` - The number of bytes written. - """ - - data = await self.read() - if isinstance(fp, io.IOBase) and fp.writable(): - written = fp.write(data) - if seek_begin: - fp.seek(0) - return written - else: - with open(fp, 'wb') as f: - return f.write(data) + return await super().read()