diff --git a/CHANGELOG.md b/CHANGELOG.md index 649b6c1901..145c298d7f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,10 @@ These changes are available on the `master` branch, but have not yet been releas ### Added +- Added `Attachment.read_chunked` and added optional `chunksize` argument to + `Attachment.save` for retrieving attachments in chunks. + ([#2956](https://github.com/Pycord-Development/pycord/pull/2956)) + ### Changed ### Fixed diff --git a/discord/http.py b/discord/http.py index 47ecfb0325..ae64703ba6 100644 --- a/discord/http.py +++ b/discord/http.py @@ -29,7 +29,15 @@ import logging import sys import weakref -from typing import TYPE_CHECKING, Any, Coroutine, Iterable, Sequence, TypeVar +from typing import ( + TYPE_CHECKING, + Any, + AsyncGenerator, + Coroutine, + Iterable, + Sequence, + TypeVar, +) from urllib.parse import quote as _uriquote import aiohttp @@ -406,6 +414,21 @@ async def get_from_cdn(self, url: str) -> bytes: else: raise HTTPException(resp, "failed to get asset") + async def stream_from_cdn(self, url: str, chunksize: int) -> AsyncGenerator[bytes]: + if not isinstance(chunksize, int) or chunksize < 1: + raise InvalidArgument("The chunksize must be a positive integer.") + + async with self.__session.get(url) as resp: + if resp.status == 200: + async for chunk in resp.content.iter_chunked(chunksize): + yield chunk + elif resp.status == 404: + raise NotFound(resp, "asset not found") + elif resp.status == 403: + raise Forbidden(resp, "cannot retrieve asset") + else: + raise HTTPException(resp, "failed to get asset") + # state management async def close(self) -> None: diff --git a/discord/message.py b/discord/message.py index 117050d2d3..ad2e986ec8 100644 --- a/discord/message.py +++ b/discord/message.py @@ -32,6 +32,7 @@ from typing import ( TYPE_CHECKING, Any, + AsyncGenerator, Callable, ClassVar, Sequence, @@ -290,6 +291,7 @@ async def save( *, seek_begin: bool = True, use_cached: bool = False, + chunksize: int | None = None, ) -> int: """|coro| @@ -311,6 +313,8 @@ async def save( after the message is deleted. Note that this can still fail to download deleted attachments if too much time has passed, and it does not work on some types of attachments. + chunksize: Optional[:class:`int`] + The maximum size of each chunk to process. Returns ------- @@ -323,16 +327,33 @@ async def save( Saving the attachment failed. NotFound The attachment was deleted. + InvalidArgument + Argument `chunksize` is less than 1. """ - data = await self.read(use_cached=use_cached) + if chunksize is not None: + data = self.read_chunked(use_cached=use_cached, chunksize=chunksize) + else: + data = await self.read(use_cached=use_cached) + if isinstance(fp, io.BufferedIOBase): - written = fp.write(data) + if chunksize: + written = 0 + async for chunk in data: + written += fp.write(chunk) + else: + written = fp.write(data) if seek_begin: fp.seek(0) return written else: with open(fp, "wb") as f: - return f.write(data) + if chunksize: + written = 0 + async for chunk in data: + written += f.write(chunk) + return written + else: + return f.write(data) async def read(self, *, use_cached: bool = False) -> bytes: """|coro| @@ -369,6 +390,45 @@ async def read(self, *, use_cached: bool = False) -> bytes: data = await self._http.get_from_cdn(url) return data + async def read_chunked( + self, chunksize: int, *, use_cached: bool = False + ) -> AsyncGenerator[bytes]: + """|coro| + + Retrieves the content of this attachment in chunks as a :class:`AsyncGenerator` object of bytes. + + Parameters + ---------- + chunksize: :class:`int` + The maximum size of each chunk to process. + use_cached: :class:`bool` + Whether to use :attr:`proxy_url` rather than :attr:`url` when downloading + the attachment. This will allow attachments to be saved after deletion + more often, compared to the regular URL which is generally deleted right + after the message is deleted. Note that this can still fail to download + deleted attachments if too much time has passed, and it does not work + on some types of attachments. + + Yields + ------ + :class:`bytes` + A chunk of the file. + + Raises + ------ + HTTPException + Downloading the attachment failed. + Forbidden + You do not have permissions to access this attachment + NotFound + The attachment was deleted. + InvalidArgument + Argument `chunksize` is less than 1. + """ + url = self.proxy_url if use_cached else self.url + async for chunk in self._http.stream_from_cdn(url, chunksize): + yield chunk + async def to_file(self, *, use_cached: bool = False, spoiler: bool = False) -> File: """|coro|