feat: make attachment downloading chunkable

CHANGELOG.md

@@ -30,6 +30,9 @@ These changes are available on the `master` branch, but have not yet been releas
     `ui.MentionableSelect`, and `ui.ChannelSelect`.
 - Added `ui.FileUpload` for modals and the `FileUpload` component.
   ([#2938](https://github.com/Pycord-Development/pycord/pull/2938))
+- Added `Attachment.read_chunked` and optional `chunksize` argument to `Attachment.save`
+  for processing attachments in chunks.
+  ([#2956](https://github.com/Pycord-Development/pycord/pull/2956))
 
 ### Changed
 

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,23 @@ async def get_from_cdn(self, url: str) -> bytes:
             else:
                 raise HTTPException(resp, "failed to get asset")
 
+    async def get_from_cdn_stream(
+        self, url: str, chunksize: int
+    ) -> AsyncGenerator[bytes]:
+        if chunksize < 1:
+            raise InvalidArgument("The chunksize must be a positive number.")
+
+        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:
+            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, *, use_cached: bool = False, chunksize: int | None = None
+    ) -> AsyncGenerator[bytes]:
+        """|coro|
+
+        Retrieves the content of this attachment in chunks as a :class:`AsyncGenerator` object of bytes.
+
+        Parameters
+        ----------
+        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.
+        chunksize: class:`int`
+            The maximum size of each chunk to process.
+
+        Returns
+        -------
+        :class:`AsyncGenerator`
+            Generator of the contents of the attachment.
+
+        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.get_from_cdn_stream(url, chunksize):
+            yield chunk
+
     async def to_file(self, *, use_cached: bool = False, spoiler: bool = False) -> File:
         """|coro|