nonebot2/nonebot/drivers/httpx.py
2024-10-22 10:33:48 +08:00

161 lines
4.1 KiB
Python

"""[HTTPX](https://www.python-httpx.org/) 驱动适配
```bash
nb driver install httpx
# 或者
pip install nonebot2[httpx]
```
:::tip 提示
本驱动仅支持客户端 HTTP 连接
:::
FrontMatter:
mdx:
format: md
sidebar_position: 3
description: nonebot.drivers.httpx 模块
"""
from typing_extensions import override
from typing import TYPE_CHECKING, Union, Optional
from multidict import CIMultiDict
from nonebot.drivers.none import Driver as NoneDriver
from nonebot.internal.driver import Cookies, QueryTypes, CookieTypes, HeaderTypes
from nonebot.drivers import (
URL,
Request,
Response,
HTTPVersion,
HTTPClientMixin,
HTTPClientSession,
combine_driver,
)
try:
import httpx
except ModuleNotFoundError as e: # pragma: no cover
raise ImportError(
"Please install httpx first to use this driver. "
"Install with pip: `pip install nonebot2[httpx]`"
) from e
class Session(HTTPClientSession):
@override
def __init__(
self,
params: QueryTypes = None,
headers: HeaderTypes = None,
cookies: CookieTypes = None,
version: Union[str, HTTPVersion] = HTTPVersion.H11,
timeout: Optional[float] = None,
proxy: Optional[str] = None,
):
self._client: Optional[httpx.AsyncClient] = 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)
self._timeout = timeout
self._proxy = proxy
@property
def client(self) -> httpx.AsyncClient:
if self._client is None:
raise RuntimeError("Session is not initialized")
return self._client
@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,
headers=tuple(setup.headers.items()),
cookies=setup.cookies.jar,
timeout=setup.timeout,
)
return Response(
response.status_code,
headers=response.headers.multi_items(),
content=response.content,
request=setup,
)
@override
async def setup(self) -> None:
if self._client is not None:
raise RuntimeError("Session has already been initialized")
self._client = httpx.AsyncClient(
params=self._params,
headers=self._headers,
cookies=self._cookies.jar,
http2=self._version == HTTPVersion.H2,
proxies=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):
"""HTTPX Mixin"""
@property
@override
def type(self) -> str:
return "httpx"
@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
def get_session(
self,
params: QueryTypes = None,
headers: HeaderTypes = None,
cookies: CookieTypes = None,
version: Union[str, HTTPVersion] = HTTPVersion.H11,
timeout: Optional[float] = None,
proxy: Optional[str] = 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)
"""HTTPX Driver"""