✨ Feature: 添加 httpx2 驱动器支持 (#4194)

nonebot/drivers/httpx2.py

@@ -0,0 +1,262 @@
+"""[HTTPX2](https://github.com/pydantic/httpx2) 驱动适配
+
+```bash
+nb driver install httpx2
+# 或者
+pip install nonebot2[httpx2]
+```
+
+:::tip[提示]
+本驱动仅支持客户端 HTTP 连接
+:::
+
+FrontMatter:
+    mdx:
+        format: md
+    sidebar_position: 7
+    description: nonebot.drivers.httpx2 模块
+"""
+
+from collections.abc import AsyncGenerator
+from typing import TYPE_CHECKING
+from typing_extensions import override
+
+from multidict import CIMultiDict
+
+from nonebot.drivers import (
+    URL,
+    HTTPClientMixin,
+    HTTPClientSession,
+    HTTPVersion,
+    Request,
+    Response,
+    combine_driver,
+)
+from nonebot.drivers.none import Driver as NoneDriver
+from nonebot.internal.driver import (
+    DEFAULT_TIMEOUT,
+    Cookies,
+    CookieTypes,
+    HeaderTypes,
+    QueryTypes,
+    Timeout,
+    TimeoutTypes,
+)
+from nonebot.utils import UNSET, UnsetType, exclude_unset
+
+try:
+    import httpx2
+except ModuleNotFoundError as e:  # pragma: no cover
+    raise ImportError(
+        "Please install httpx2 first to use this driver. "
+        "Install with pip: `pip install nonebot2[httpx2]`"
+    ) from e
+
+
+class Session(HTTPClientSession):
+    @override
+    def __init__(
+        self,
+        params: QueryTypes = None,
+        headers: HeaderTypes = None,
+        cookies: CookieTypes = None,
+        version: str | HTTPVersion = HTTPVersion.H11,
+        timeout: TimeoutTypes | UnsetType = UNSET,
+        proxy: str | None = None,
+    ):
+        self._client: httpx2.AsyncClient | None = None
+
+        self._params = (
+            tuple(URL.build(query=params).query.items()) if params is not None else None
+        )
+        self._headers = (
+            tuple(CIMultiDict(headers).items()) if headers is not None else None
+        )
+        self._cookies = Cookies(cookies)
+        self._version = HTTPVersion(version)
+
+        _timeout = None
+        if isinstance(timeout, Timeout):
+            avg_timeout = timeout.total and timeout.total / 4
+            timeout_kwargs: dict[str, float | None] = exclude_unset(
+                {
+                    "timeout": avg_timeout,
+                    "connect": timeout.connect,
+                    "read": timeout.read,
+                }
+            )
+            if timeout_kwargs:
+                _timeout = httpx2.Timeout(**timeout_kwargs)
+        elif timeout is not UNSET:
+            _timeout = httpx2.Timeout(timeout)
+
+        if _timeout is None:
+            avg_timeout = DEFAULT_TIMEOUT.total and DEFAULT_TIMEOUT.total / 4
+            _timeout = httpx2.Timeout(
+                **exclude_unset(
+                    {
+                        "timeout": avg_timeout,
+                        "connect": DEFAULT_TIMEOUT.connect,
+                        "read": DEFAULT_TIMEOUT.read,
+                    }
+                )
+            )
+
+        self._timeout = _timeout
+        self._proxy = proxy
+
+    @property
+    def client(self) -> httpx2.AsyncClient:
+        if self._client is None:
+            raise RuntimeError("Session is not initialized")
+        return self._client
+
+    def _get_timeout(self, timeout: TimeoutTypes | UnsetType) -> httpx2.Timeout:
+        _timeout = None
+        if isinstance(timeout, Timeout):
+            avg_timeout = timeout.total and timeout.total / 4
+            timeout_kwargs: dict[str, float | None] = exclude_unset(
+                {
+                    "timeout": avg_timeout,
+                    "connect": timeout.connect,
+                    "read": timeout.read,
+                }
+            )
+            if timeout_kwargs:
+                _timeout = httpx2.Timeout(**timeout_kwargs)
+        elif timeout is not UNSET:
+            _timeout = httpx2.Timeout(timeout)
+
+        if _timeout is None:
+            return self._timeout
+        return _timeout
+
+    @override
+    async def request(self, setup: Request) -> Response:
+        response = await self.client.request(
+            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=self._get_timeout(setup.timeout),
+        )
+        return Response(
+            response.status_code,
+            headers=response.headers.multi_items(),
+            content=response.content,
+            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=self._get_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:
+            raise RuntimeError("Session has already been initialized")
+        self._client = httpx2.AsyncClient(
+            params=self._params,
+            headers=self._headers,
+            cookies=self._cookies.jar,
+            http2=self._version == HTTPVersion.H2,
+            proxy=self._proxy,
+            follow_redirects=True,
+        )
+        await self._client.__aenter__()
+
+    @override
+    async def close(self) -> None:
+        try:
+            if self._client is not None:
+                await self._client.aclose()
+        finally:
+            self._client = None
+
+
+class Mixin(HTTPClientMixin):
+    """HTTPX2 Mixin"""
+
+    @property
+    @override
+    def type(self) -> str:
+        return "httpx2"
+
+    @override
+    async def request(self, setup: Request) -> Response:
+        async with self.get_session(
+            version=setup.version, proxy=setup.proxy
+        ) 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,
+        params: QueryTypes = None,
+        headers: HeaderTypes = None,
+        cookies: CookieTypes = None,
+        version: str | HTTPVersion = HTTPVersion.H11,
+        timeout: TimeoutTypes = None,
+        proxy: str | None = None,
+    ) -> Session:
+        return Session(
+            params=params,
+            headers=headers,
+            cookies=cookies,
+            version=version,
+            timeout=timeout,
+            proxy=proxy,
+        )
+
+
+if TYPE_CHECKING:
+
+    class Driver(Mixin, NoneDriver): ...
+
+else:
+    Driver = combine_driver(NoneDriver, Mixin)
+    """HTTPX2 Driver"""