feat: make attachment downloading chunkable (#2956)

Co-authored-by: Copilot <[email protected]>
Co-authored-by: plun1331 <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Paillat <[email protected]>
Co-authored-by: Lala Sabathil <[email protected]>
Co-authored-by: Paillat <[email protected]>
Co-authored-by: Lala Sabathil <[email protected]>

Author: 6 people

CHANGELOG.md

@@ -14,6 +14,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

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:

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|