Skip to content

Commit

Permalink
Refactor save() and read() into AssetMixin
Browse files Browse the repository at this point in the history
  • Loading branch information
@Rapptz
Rapptz committed Apr 17, 2021
1 parent f6fcffb commit fed259a
Show file tree
Hide file tree
Showing 4 changed files with 78 additions and 206 deletions.
149 changes: 71 additions & 78 deletions discord/asset.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@

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
Expand All @@ -45,7 +45,76 @@
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
Expand Down Expand Up @@ -153,19 +222,6 @@ def _from_sticker(cls, state, sticker_id: int, sticker_hash: str) -> 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

Expand Down Expand Up @@ -332,66 +388,3 @@ def with_static_format(self, format: ValidStaticFormatTypes) -> 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)
61 changes: 2 additions & 59 deletions discord/emoji.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@

import io

from .asset import Asset
from .asset import Asset, AssetMixin
from . import utils
from .partial_emoji import _EmojiTag
from .user import User
Expand All @@ -33,7 +33,7 @@
'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
Expand Down Expand Up @@ -221,60 +221,3 @@ async def edit(self, *, name=None, roles=None, reason=None):
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)
2 changes: 1 addition & 1 deletion discord/message.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -172,7 +172,7 @@ async def save(self, fp, *, seek_begin=True, use_cached=False):
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)
Expand Down
72 changes: 4 additions & 68 deletions discord/partial_emoji.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@

import io

from .asset import Asset
from .asset import Asset, AssetMixin
from .errors import DiscordException, InvalidArgument
from . import utils

Expand All @@ -35,7 +35,7 @@
class _EmojiTag:
__slots__ = ()

class PartialEmoji(_EmojiTag):
class PartialEmoji(_EmojiTag, AssetMixin):
"""Represents a "partial" emoji.
This model will be given in two scenarios:
Expand Down Expand Up @@ -163,72 +163,8 @@ def url(self) -> str:
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()

0 comments on commit fed259a

Please sign in to comment.