nonebot2/nonebot/drivers/httpx.py
2021-12-24 23:15:28 +08:00

52 lines
1.4 KiB
Python

from nonebot.typing import overrides
from nonebot.drivers._block_driver import BlockDriver
from nonebot.drivers import (
Request,
Response,
WebSocket,
HTTPVersion,
ForwardMixin,
combine_driver,
)
try:
import httpx
except ImportError:
raise ImportError("Please install httpx by using `pip install nonebot2[httpx]`")
class Mixin(ForwardMixin):
@property
@overrides(ForwardMixin)
def type(self) -> str:
return "httpx"
@overrides(ForwardMixin)
async def request(self, setup: Request) -> Response:
async with httpx.AsyncClient(
http2=setup.version == HTTPVersion.H2, follow_redirects=True
) as client:
response = await client.request(
setup.method,
str(setup.url),
content=setup.content,
data=setup.data,
json=setup.json,
files=setup.files,
headers=tuple(setup.headers.items()),
timeout=setup.timeout,
)
return Response(
response.status_code,
headers=response.headers,
content=response.content,
request=setup,
)
@overrides(ForwardMixin)
async def websocket(self, setup: Request) -> WebSocket:
return await super(Mixin, self).websocket(setup)
Driver = combine_driver(BlockDriver, Mixin)