✨ Feature: 为 HTTPClient 等内置驱动器添加流式请求方法 (#3560)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

nonebot/drivers/aiohttp.py

@@ -125,6 +125,51 @@ class Session(HTTPClientSession):
                 request=setup,
             )
 
+    @override
+    async def stream_request(
+        self,
+        setup: Request,
+        *,
+        chunk_size: int = 1024,
+    ) -> AsyncGenerator[Response, None]:
+        if self._params:
+            url = setup.url.with_query({**self._params, **setup.url.query})
+        else:
+            url = setup.url
+
+        data = setup.data
+        if setup.files:
+            data = aiohttp.FormData(data or {}, quote_fields=False)
+            for name, file in setup.files:
+                data.add_field(name, file[1], content_type=file[2], filename=file[0])
+
+        cookies = (
+            (cookie.name, cookie.value)
+            for cookie in setup.cookies
+            if cookie.value is not None
+        )
+
+        timeout = aiohttp.ClientTimeout(setup.timeout)
+
+        async with self.client.request(
+            setup.method,
+            url,
+            data=setup.content or data,
+            json=setup.json,
+            cookies=cookies,
+            headers=setup.headers,
+            proxy=setup.proxy or self._proxy,
+            timeout=timeout,
+        ) as response:
+            response_headers = response.headers.copy()
+            async for chunk in response.content.iter_chunked(chunk_size):
+                yield Response(
+                    response.status,
+                    headers=response_headers,
+                    content=chunk,
+                    request=setup,
+                )
+
     @override
     async def setup(self) -> None:
         if self._client is not None:
@@ -160,6 +205,17 @@ class Mixin(HTTPClientMixin, WebSocketClientMixin):
         async with self.get_session() as session:
             return await session.request(setup)
 
+    @override
+    async def stream_request(
+        self,
+        setup: Request,
+        *,
+        chunk_size: int = 1024,
+    ) -> AsyncGenerator[Response, None]:
+        async with self.get_session() as session:
+            async for response in session.stream_request(setup, chunk_size=chunk_size):
+                yield response
+
     @override
     @asynccontextmanager
     async def websocket(self, setup: Request) -> AsyncGenerator["WebSocket", None]:

nonebot/drivers/httpx.py

@@ -17,6 +17,7 @@ FrontMatter:
     description: nonebot.drivers.httpx 模块
 """
 
+from collections.abc import AsyncGenerator
 from typing import TYPE_CHECKING, Optional, Union
 from typing_extensions import override
 
@@ -95,6 +96,35 @@ class Session(HTTPClientSession):
             request=setup,
         )
 
+    @override
+    async def stream_request(
+        self,
+        setup: Request,
+        *,
+        chunk_size: int = 1024,
+    ) -> AsyncGenerator[Response, None]:
+        async with self.client.stream(
+            setup.method,
+            str(setup.url),
+            content=setup.content,
+            data=setup.data,
+            files=setup.files,
+            json=setup.json,
+            # ensure the params priority
+            params=setup.url.raw_query_string,
+            headers=tuple(setup.headers.items()),
+            cookies=setup.cookies.jar,
+            timeout=setup.timeout,
+        ) as response:
+            response_headers = response.headers.multi_items()
+            async for chunk in response.aiter_bytes(chunk_size=chunk_size):
+                yield Response(
+                    response.status_code,
+                    headers=response_headers,
+                    content=chunk,
+                    request=setup,
+                )
+
     @override
     async def setup(self) -> None:
         if self._client is not None:
@@ -133,6 +163,19 @@ class Mixin(HTTPClientMixin):
         ) as session:
             return await session.request(setup)
 
+    @override
+    async def stream_request(
+        self,
+        setup: Request,
+        *,
+        chunk_size: int = 1024,
+    ) -> AsyncGenerator[Response, None]:
+        async with self.get_session(
+            version=setup.version, proxy=setup.proxy
+        ) as session:
+            async for response in session.stream_request(setup, chunk_size=chunk_size):
+                yield response
+
     @override
     def get_session(
         self,

nonebot/internal/driver/abstract.py

@@ -255,6 +255,17 @@ class HTTPClientSession(abc.ABC):
         """发送一个 HTTP 请求"""
         raise NotImplementedError
 
+    @abc.abstractmethod
+    async def stream_request(
+        self,
+        setup: Request,
+        *,
+        chunk_size: int = 1024,
+    ) -> AsyncGenerator[Response, None]:
+        """发送一个 HTTP 流式请求"""
+        raise NotImplementedError
+        yield  # used for static type checking's generator detection
+
     @abc.abstractmethod
     async def setup(self) -> None:
         """初始化会话"""
@@ -286,6 +297,17 @@ class HTTPClientMixin(ForwardMixin):
         """发送一个 HTTP 请求"""
         raise NotImplementedError
 
+    @abc.abstractmethod
+    async def stream_request(
+        self,
+        setup: Request,
+        *,
+        chunk_size: int = 1024,
+    ) -> AsyncGenerator[Response, None]:
+        """发送一个 HTTP 流式请求"""
+        raise NotImplementedError
+        yield  # used for static type checking's generator detection
+
     @abc.abstractmethod
     def get_session(
         self,