ProjectHentai/nonebot2
6
0
Fork 0
mirror of https://github.com/nonebot/nonebot2.git synced 2025-09-07 04:26:45 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

🎨 format code using black and isort

Browse Source
This commit is contained in:
yanyongyu
2021-11-22 23:21:26 +08:00
parent 602185a34e
commit a98d98cd12
86 changed files with 2893 additions and 2095 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
nonebot/drivers/__init__.py
View File

@ -8,8 +8,17 @@
import abc
import asyncio
from dataclasses import field, dataclass
from typing import (TYPE_CHECKING, Any, Set, Dict, Type, Union, Callable,
Optional, Awaitable)
from typing import (
TYPE_CHECKING,
Any,
Set,
Dict,
Type,
Union,
Callable,
Optional,
Awaitable,
)
from nonebot.log import logger
from nonebot.utils import escape_tag
@ -90,12 +99,14 @@ class Driver(abc.ABC):
"""
if name in self._adapters:
logger.opt(colors=True).debug(
f'Adapter "<y>{escape_tag(name)}</y>" already exists')
f'Adapter "<y>{escape_tag(name)}</y>" already exists'
)
return
self._adapters[name] = adapter
adapter.register(self, self.config, **kwargs)
logger.opt(colors=True).debug(
f'Succeeded to load adapter "<y>{escape_tag(name)}</y>"')
f'Succeeded to load adapter "<y>{escape_tag(name)}</y>"'
)
@property
@abc.abstractmethod
@ -121,7 +132,8 @@ class Driver(abc.ABC):
* ``**kwargs``
"""
logger.opt(colors=True).debug(
f"<g>Loaded adapters: {escape_tag(', '.join(self._adapters))}</g>")
f"<g>Loaded adapters: {escape_tag(', '.join(self._adapters))}</g>"
)
@abc.abstractmethod
def on_startup(self, func: Callable) -> Callable:
@ -146,8 +158,7 @@ class Driver(abc.ABC):
self._bot_connection_hook.add(func)
return func
def on_bot_disconnect(
self, func: T_BotDisconnectionHook) -> T_BotDisconnectionHook:
def on_bot_disconnect(self, func: T_BotDisconnectionHook) -> T_BotDisconnectionHook:
"""
:说明:
@ -172,7 +183,8 @@ class Driver(abc.ABC):
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error when running WebSocketConnection hook. "
"Running cancelled!</bg #f8bbd0></r>")
"Running cancelled!</bg #f8bbd0></r>"
)
asyncio.create_task(_run_hook(bot))
@ -189,7 +201,8 @@ class Driver(abc.ABC):
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error when running WebSocketDisConnection hook. "
"Running cancelled!</bg #f8bbd0></r>")
"Running cancelled!</bg #f8bbd0></r>"
)
asyncio.create_task(_run_hook(bot))
@ -201,8 +214,8 @@ class ForwardDriver(Driver):
@abc.abstractmethod
def setup_http_polling(
self, setup: Union["HTTPPollingSetup",
Callable[[], Awaitable["HTTPPollingSetup"]]]
self,
setup: Union["HTTPPollingSetup", Callable[[], Awaitable["HTTPPollingSetup"]]],
) -> None:
"""
:说明:
@ -217,8 +230,7 @@ class ForwardDriver(Driver):
@abc.abstractmethod
def setup_websocket(
self, setup: Union["WebSocketSetup",
Callable[[], Awaitable["WebSocketSetup"]]]
self, setup: Union["WebSocketSetup", Callable[[], Awaitable["WebSocketSetup"]]]
) -> None:
"""
:说明:
@ -288,6 +300,7 @@ class HTTPRequest(HTTPConnection):
.. _asgi http scope:
https://asgi.readthedocs.io/en/latest/specs/www.html#http-connection-scope
"""
method: str = "GET"
"""The HTTP method name, uppercased."""
body: bytes = b""
@ -309,6 +322,7 @@ class HTTPResponse:
.. _asgi http scope:
https://asgi.readthedocs.io/en/latest/specs/www.html#http-connection-scope
"""
status: int
"""HTTP status code."""
body: Optional[bytes] = None
@ -416,5 +430,5 @@ class WebSocketSetup:
"""URL"""
headers: Dict[str, str] = field(default_factory=dict)
"""HTTP headers"""
reconnect_interval: float = 3.
reconnect_interval: float = 3.0
"""WebSocket 重连间隔"""

112
nonebot/drivers/aiohttp.py
View File

@ -20,13 +20,16 @@ from nonebot.typing import overrides
from nonebot.utils import escape_tag
from nonebot.config import Env, Config
from nonebot.drivers import WebSocket as BaseWebSocket
from nonebot.drivers import (HTTPRequest, ForwardDriver, WebSocketSetup,
HTTPPollingSetup)
from nonebot.drivers import (
HTTPRequest,
ForwardDriver,
WebSocketSetup,
HTTPPollingSetup,
)
STARTUP_FUNC = Callable[[], Awaitable[None]]
SHUTDOWN_FUNC = Callable[[], Awaitable[None]]
HTTPPOLLING_SETUP = Union[HTTPPollingSetup,
Callable[[], Awaitable[HTTPPollingSetup]]]
HTTPPOLLING_SETUP = Union[HTTPPollingSetup, Callable[[], Awaitable[HTTPPollingSetup]]]
WEBSOCKET_SETUP = Union[WebSocketSetup, Callable[[], Awaitable[WebSocketSetup]]]
HANDLED_SIGNALS = (
signal.SIGINT, # Unix signal 2. Sent by Ctrl+C.
@ -146,7 +149,8 @@ class Driver(ForwardDriver):
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error when running startup function. "
"Ignored!</bg #f8bbd0></r>")
"Ignored!</bg #f8bbd0></r>"
)
async def main_loop(self):
await self.should_exit.wait()
@ -160,24 +164,20 @@ class Driver(ForwardDriver):
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error when running shutdown function. "
"Ignored!</bg #f8bbd0></r>")
"Ignored!</bg #f8bbd0></r>"
)
for task in self.connections:
if not task.done():
task.cancel()
await asyncio.sleep(0.1)
tasks = [
t for t in asyncio.all_tasks() if t is not asyncio.current_task()
]
tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
if tasks and not self.force_exit:
logger.info("Waiting for tasks to finish. (CTRL+C to force quit)")
while tasks and not self.force_exit:
await asyncio.sleep(0.1)
tasks = [
t for t in asyncio.all_tasks()
if t is not asyncio.current_task()
]
tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
for task in tasks:
task.cancel()
@ -209,9 +209,7 @@ class Driver(ForwardDriver):
self.should_exit.set()
async def _http_loop(self, setup: HTTPPOLLING_SETUP):
async def _build_request(
setup: HTTPPollingSetup) -> Optional[HTTPRequest]:
async def _build_request(setup: HTTPPollingSetup) -> Optional[HTTPRequest]:
url = URL(setup.url)
if not url.is_absolute() or not url.host:
logger.opt(colors=True).error(
@ -219,10 +217,15 @@ class Driver(ForwardDriver):
)
return
host = f"{url.host}:{url.port}" if url.port else url.host
return HTTPRequest(setup.http_version, url.scheme, url.path,
url.raw_query_string.encode("latin-1"), {
**setup.headers, "host": host
}, setup.method, setup.body)
return HTTPRequest(
setup.http_version,
url.scheme,
url.path,
url.raw_query_string.encode("latin-1"),
{**setup.headers, "host": host},
setup.method,
setup.body,
)
bot: Optional[Bot] = None
request: Optional[HTTPRequest] = None
@ -230,7 +233,8 @@ class Driver(ForwardDriver):
logger.opt(colors=True).info(
f"Start http polling for <y>{escape_tag(setup.adapter.upper())} "
f"Bot {escape_tag(setup.self_id)}</y>")
f"Bot {escape_tag(setup.self_id)}</y>"
)
try:
async with aiohttp.ClientSession() as session:
@ -244,7 +248,8 @@ class Driver(ForwardDriver):
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error while parsing setup "
f"{escape_tag(repr(setup))}.</bg #f8bbd0></r>")
f"{escape_tag(repr(setup))}.</bg #f8bbd0></r>"
)
await asyncio.sleep(3)
continue
@ -286,19 +291,22 @@ class Driver(ForwardDriver):
)
try:
async with session.request(request.method,
setup_.url,
data=request.body,
headers=headers,
timeout=timeout,
version=version) as response:
async with session.request(
request.method,
setup_.url,
data=request.body,
headers=headers,
timeout=timeout,
version=version,
) as response:
response.raise_for_status()
data = await response.read()
asyncio.create_task(bot.handle_message(data))
except aiohttp.ClientResponseError as e:
logger.opt(colors=True, exception=e).error(
f"<r><bg #f8bbd0>Error occurred while requesting {escape_tag(setup_.url)}. "
"Try to reconnect...</bg #f8bbd0></r>")
"Try to reconnect...</bg #f8bbd0></r>"
)
await asyncio.sleep(setup_.poll_interval)
@ -307,7 +315,8 @@ class Driver(ForwardDriver):
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Unexpected exception occurred "
"while http polling</bg #f8bbd0></r>")
"while http polling</bg #f8bbd0></r>"
)
finally:
if bot:
self._bot_disconnect(bot)
@ -327,7 +336,8 @@ class Driver(ForwardDriver):
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error while parsing setup "
f"{escape_tag(repr(setup))}.</bg #f8bbd0></r>")
f"{escape_tag(repr(setup))}.</bg #f8bbd0></r>"
)
await asyncio.sleep(3)
continue
@ -346,17 +356,21 @@ class Driver(ForwardDriver):
f"Bot {setup_.self_id} from adapter {setup_.adapter} connecting to {url}"
)
try:
async with session.ws_connect(url,
headers=headers,
timeout=30.) as ws:
async with session.ws_connect(
url, headers=headers, timeout=30.0
) as ws:
logger.opt(colors=True).info(
f"WebSocket Connection to <y>{escape_tag(setup_.adapter.upper())} "
f"Bot {escape_tag(setup_.self_id)}</y> succeeded!"
)
request = WebSocket(
"1.1", url.scheme, url.path,
url.raw_query_string.encode("latin-1"), headers,
ws)
"1.1",
url.scheme,
url.path,
url.raw_query_string.encode("latin-1"),
headers,
ws,
)
BotClass = self._adapters[setup_.adapter]
bot = BotClass(setup_.self_id, request)
@ -365,25 +379,30 @@ class Driver(ForwardDriver):
msg = await ws.receive()
if msg.type == aiohttp.WSMsgType.text:
asyncio.create_task(
bot.handle_message(msg.data.encode()))
bot.handle_message(msg.data.encode())
)
elif msg.type == aiohttp.WSMsgType.binary:
asyncio.create_task(
bot.handle_message(msg.data))
asyncio.create_task(bot.handle_message(msg.data))
elif msg.type == aiohttp.WSMsgType.error:
logger.opt(colors=True).error(
"<r><bg #f8bbd0>Error while handling websocket frame. "
"Try to reconnect...</bg #f8bbd0></r>")
"Try to reconnect...</bg #f8bbd0></r>"
)
break
else:
logger.opt(colors=True).error(
"<r><bg #f8bbd0>WebSocket connection closed by peer. "
"Try to reconnect...</bg #f8bbd0></r>")
"Try to reconnect...</bg #f8bbd0></r>"
)
break
except (aiohttp.ClientResponseError,
aiohttp.ClientConnectionError) as e:
except (
aiohttp.ClientResponseError,
aiohttp.ClientConnectionError,
) as e:
logger.opt(colors=True, exception=e).error(
f"<r><bg #f8bbd0>Error while connecting to {escape_tag(str(url))}. "
"Try to reconnect...</bg #f8bbd0></r>")
"Try to reconnect...</bg #f8bbd0></r>"
)
finally:
if bot:
self._bot_disconnect(bot)
@ -395,7 +414,8 @@ class Driver(ForwardDriver):
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Unexpected exception occurred "
"while websocket loop</bg #f8bbd0></r>")
"while websocket loop</bg #f8bbd0></r>"
)
@dataclass

165
nonebot/drivers/fastapi.py
View File

@ -32,11 +32,15 @@ from nonebot.typing import overrides
from nonebot.utils import escape_tag
from nonebot.config import Config as NoneBotConfig
from nonebot.drivers import WebSocket as BaseWebSocket
from nonebot.drivers import (HTTPRequest, ForwardDriver, ReverseDriver,
WebSocketSetup, HTTPPollingSetup)
from nonebot.drivers import (
HTTPRequest,
ForwardDriver,
ReverseDriver,
WebSocketSetup,
HTTPPollingSetup,
)
HTTPPOLLING_SETUP = Union[HTTPPollingSetup,
Callable[[], Awaitable[HTTPPollingSetup]]]
HTTPPOLLING_SETUP = Union[HTTPPollingSetup, Callable[[], Awaitable[HTTPPollingSetup]]]
WEBSOCKET_SETUP = Union[WebSocketSetup, Callable[[], Awaitable[WebSocketSetup]]]
@ -44,6 +48,7 @@ class Config(BaseSettings):
"""
FastAPI 驱动框架设置，详情参考 FastAPI 文档
"""
fastapi_openapi_url: Optional[str] = None
"""
:类型:
@ -226,12 +231,14 @@ class Driver(ReverseDriver, ForwardDriver):
self.websockets.append(setup)
@overrides(ReverseDriver)
def run(self,
host: Optional[str] = None,
port: Optional[int] = None,
*,
app: Optional[str] = None,
**kwargs):
def run(
self,
host: Optional[str] = None,
port: Optional[int] = None,
*,
app: Optional[str] = None,
**kwargs,
):
"""使用 ``uvicorn`` 启动 FastAPI"""
super().run(host, port, app, **kwargs)
LOGGING_CONFIG = {
@ -243,10 +250,7 @@ class Driver(ReverseDriver, ForwardDriver):
},
},
"loggers": {
"uvicorn.error": {
"handlers": ["default"],
"level": "INFO"
},
"uvicorn.error": {"handlers": ["default"], "level": "INFO"},
"uvicorn.access": {
"handlers": ["default"],
"level": "INFO",
@ -258,15 +262,16 @@ class Driver(ReverseDriver, ForwardDriver):
host=host or str(self.config.host),
port=port or self.config.port,
reload=self.fastapi_config.fastapi_reload
if self.fastapi_config.fastapi_reload is not None else
(bool(app) and self.config.debug),
if self.fastapi_config.fastapi_reload is not None
else (bool(app) and self.config.debug),
reload_dirs=self.fastapi_config.fastapi_reload_dirs,
reload_delay=self.fastapi_config.fastapi_reload_delay,
reload_includes=self.fastapi_config.fastapi_reload_includes,
reload_excludes=self.fastapi_config.fastapi_reload_excludes,
debug=self.config.debug,
log_config=LOGGING_CONFIG,
**kwargs)
**kwargs,
)
def _run_forward(self):
for setup in self.http_pollings:
@ -287,39 +292,49 @@ class Driver(ReverseDriver, ForwardDriver):
logger.warning(
f"Unknown adapter {adapter}. Please register the adapter before use."
)
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND,
detail="adapter not found")
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND, detail="adapter not found"
)
# 创建 Bot 对象
BotClass = self._adapters[adapter]
http_request = HTTPRequest(request.scope["http_version"],
request.url.scheme, request.url.path,
request.scope["query_string"],
dict(request.headers), request.method, data)
x_self_id, response = await BotClass.check_permission(
self, http_request)
http_request = HTTPRequest(
request.scope["http_version"],
request.url.scheme,
request.url.path,
request.scope["query_string"],
dict(request.headers),
request.method,
data,
)
x_self_id, response = await BotClass.check_permission(self, http_request)
if not x_self_id:
raise HTTPException(
response and response.status or 401, response and
response.body and response.body.decode("utf-8"))
response and response.status or 401,
response and response.body and response.body.decode("utf-8"),
)
if x_self_id in self._clients:
logger.warning("There's already a reverse websocket connection,"
"so the event may be handled twice.")
logger.warning(
"There's already a reverse websocket connection,"
"so the event may be handled twice."
)
bot = BotClass(x_self_id, http_request)
asyncio.create_task(bot.handle_message(data))
return Response(response and response.body,
response and response.status or 200)
return Response(response and response.body, response and response.status or 200)
async def _handle_ws_reverse(self, adapter: str,
websocket: FastAPIWebSocket):
ws = WebSocket(websocket.scope.get("http_version",
"1.1"), websocket.url.scheme,
websocket.url.path, websocket.scope["query_string"],
dict(websocket.headers), websocket)
async def _handle_ws_reverse(self, adapter: str, websocket: FastAPIWebSocket):
ws = WebSocket(
websocket.scope.get("http_version", "1.1"),
websocket.url.scheme,
websocket.url.path,
websocket.scope["query_string"],
dict(websocket.headers),
websocket,
)
if adapter not in self._adapters:
logger.warning(
@ -349,7 +364,8 @@ class Driver(ReverseDriver, ForwardDriver):
await ws.accept()
logger.opt(colors=True).info(
f"WebSocket Connection from <y>{escape_tag(adapter.upper())} "
f"Bot {escape_tag(self_id)}</y> Accepted!")
f"Bot {escape_tag(self_id)}</y> Accepted!"
)
self._bot_connect(bot)
@ -362,7 +378,8 @@ class Driver(ReverseDriver, ForwardDriver):
break
except Exception as e:
logger.opt(exception=e).error(
"Error when receiving data from websocket.")
"Error when receiving data from websocket."
)
break
asyncio.create_task(bot.handle_message(data.encode()))
@ -370,9 +387,7 @@ class Driver(ReverseDriver, ForwardDriver):
self._bot_disconnect(bot)
async def _http_loop(self, setup: HTTPPOLLING_SETUP):
async def _build_request(
setup: HTTPPollingSetup) -> Optional[HTTPRequest]:
async def _build_request(setup: HTTPPollingSetup) -> Optional[HTTPRequest]:
url = httpx.URL(setup.url)
if not url.netloc:
logger.opt(colors=True).error(
@ -380,9 +395,14 @@ class Driver(ReverseDriver, ForwardDriver):
)
return
return HTTPRequest(
setup.http_version, url.scheme, url.path, url.query, {
**setup.headers, "host": url.netloc.decode("ascii")
}, setup.method, setup.body)
setup.http_version,
url.scheme,
url.path,
url.query,
{**setup.headers, "host": url.netloc.decode("ascii")},
setup.method,
setup.body,
)
bot: Optional[Bot] = None
request: Optional[HTTPRequest] = None
@ -390,11 +410,11 @@ class Driver(ReverseDriver, ForwardDriver):
logger.opt(colors=True).info(
f"Start http polling for <y>{escape_tag(setup.adapter.upper())} "
f"Bot {escape_tag(setup.self_id)}</y>")
f"Bot {escape_tag(setup.self_id)}</y>"
)
try:
async with httpx.AsyncClient(http2=True,
follow_redirects=True) as session:
async with httpx.AsyncClient(http2=True, follow_redirects=True) as session:
while not self.shutdown.is_set():
try:
@ -405,7 +425,8 @@ class Driver(ReverseDriver, ForwardDriver):
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error while parsing setup "
f"{escape_tag(repr(setup))}.</bg #f8bbd0></r>")
f"{escape_tag(repr(setup))}.</bg #f8bbd0></r>"
)
await asyncio.sleep(3)
continue
@ -432,18 +453,21 @@ class Driver(ReverseDriver, ForwardDriver):
f"Bot {setup_.self_id} from adapter {setup_.adapter} request {setup_.url}"
)
try:
response = await session.request(request.method,
setup_.url,
content=request.body,
headers=headers,
timeout=30.)
response = await session.request(
request.method,
setup_.url,
content=request.body,
headers=headers,
timeout=30.0,
)
response.raise_for_status()
data = response.read()
asyncio.create_task(bot.handle_message(data))
except httpx.HTTPError as e:
logger.opt(colors=True, exception=e).error(
f"<r><bg #f8bbd0>Error occurred while requesting {escape_tag(setup_.url)}. "
"Try to reconnect...</bg #f8bbd0></r>")
"Try to reconnect...</bg #f8bbd0></r>"
)
await asyncio.sleep(setup_.poll_interval)
@ -452,7 +476,8 @@ class Driver(ReverseDriver, ForwardDriver):
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Unexpected exception occurred "
"while http polling</bg #f8bbd0></r>")
"while http polling</bg #f8bbd0></r>"
)
finally:
if bot:
self._bot_disconnect(bot)
@ -471,7 +496,8 @@ class Driver(ReverseDriver, ForwardDriver):
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error while parsing setup "
f"{escape_tag(repr(setup))}.</bg #f8bbd0></r>")
f"{escape_tag(repr(setup))}.</bg #f8bbd0></r>"
)
await asyncio.sleep(3)
continue
@ -491,9 +517,11 @@ class Driver(ReverseDriver, ForwardDriver):
async with connection as ws:
logger.opt(colors=True).info(
f"WebSocket Connection to <y>{escape_tag(setup_.adapter.upper())} "
f"Bot {escape_tag(setup_.self_id)}</y> succeeded!")
request = WebSocket("1.1", url.scheme, url.path,
url.query, headers, ws)
f"Bot {escape_tag(setup_.self_id)}</y> succeeded!"
)
request = WebSocket(
"1.1", url.scheme, url.path, url.query, headers, ws
)
BotClass = self._adapters[setup_.adapter]
bot = BotClass(setup_.self_id, request)
@ -506,12 +534,14 @@ class Driver(ReverseDriver, ForwardDriver):
except ConnectionClosed:
logger.opt(colors=True).error(
"<r><bg #f8bbd0>WebSocket connection closed by peer. "
"Try to reconnect...</bg #f8bbd0></r>")
"Try to reconnect...</bg #f8bbd0></r>"
)
break
except Exception as e:
logger.opt(colors=True, exception=e).error(
f"<r><bg #f8bbd0>Error while connecting to {url}. "
"Try to reconnect...</bg #f8bbd0></r>")
"Try to reconnect...</bg #f8bbd0></r>"
)
finally:
if bot:
self._bot_disconnect(bot)
@ -523,21 +553,22 @@ class Driver(ReverseDriver, ForwardDriver):
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Unexpected exception occurred "
"while websocket loop</bg #f8bbd0></r>")
"while websocket loop</bg #f8bbd0></r>"
)
@dataclass
class WebSocket(BaseWebSocket):
websocket: Union[FastAPIWebSocket,
WebSocketClientProtocol] = None # type: ignore
websocket: Union[FastAPIWebSocket, WebSocketClientProtocol] = None # type: ignore
@property
@overrides(BaseWebSocket)
def closed(self) -> bool:
if isinstance(self.websocket, FastAPIWebSocket):
return (
self.websocket.client_state == WebSocketState.DISCONNECTED or
self.websocket.application_state == WebSocketState.DISCONNECTED)
self.websocket.client_state == WebSocketState.DISCONNECTED
or self.websocket.application_state == WebSocketState.DISCONNECTED
)
else:
return self.websocket.closed

96
nonebot/drivers/quart.py
View File

@ -30,8 +30,7 @@ try:
from quart import Quart, Request, Response
from quart import Websocket as QuartWebSocket
except ImportError:
raise ValueError(
'Please install Quart by using `pip install nonebot2[quart]`')
raise ValueError("Please install Quart by using `pip install nonebot2[quart]`")
_AsyncCallable = TypeVar("_AsyncCallable", bound=Callable[..., Coroutine])
@ -40,6 +39,7 @@ class Config(BaseSettings):
"""
Quart 驱动框架设置
"""
quart_reload: Optional[bool] = None
"""
:类型:
@ -111,11 +111,12 @@ class Driver(ReverseDriver):
self.quart_config = Config(**config.dict())
self._server_app = Quart(self.__class__.__qualname__)
self._server_app.add_url_rule("/<adapter>/http",
methods=["POST"],
view_func=self._handle_http)
self._server_app.add_websocket("/<adapter>/ws",
view_func=self._handle_ws_reverse)
self._server_app.add_url_rule(
"/<adapter>/http", methods=["POST"], view_func=self._handle_http
)
self._server_app.add_websocket(
"/<adapter>/ws", view_func=self._handle_ws_reverse
)
@property
@overrides(ReverseDriver)
@ -156,12 +157,14 @@ class Driver(ReverseDriver):
return self.server_app.after_serving(func) # type: ignore
@overrides(ReverseDriver)
def run(self,
host: Optional[str] = None,
port: Optional[int] = None,
*,
app: Optional[str] = None,
**kwargs):
def run(
self,
host: Optional[str] = None,
port: Optional[int] = None,
*,
app: Optional[str] = None,
**kwargs,
):
"""使用 ``uvicorn`` 启动 Quart"""
super().run(host, port, app, **kwargs)
LOGGING_CONFIG = {
@ -173,10 +176,7 @@ class Driver(ReverseDriver):
},
},
"loggers": {
"uvicorn.error": {
"handlers": ["default"],
"level": "INFO"
},
"uvicorn.error": {"handlers": ["default"], "level": "INFO"},
"uvicorn.access": {
"handlers": ["default"],
"level": "INFO",
@ -188,52 +188,69 @@ class Driver(ReverseDriver):
host=host or str(self.config.host),
port=port or self.config.port,
reload=self.quart_config.quart_reload
if self.quart_config.quart_reload is not None else
(bool(app) and self.config.debug),
if self.quart_config.quart_reload is not None
else (bool(app) and self.config.debug),
reload_dirs=self.quart_config.quart_reload_dirs,
reload_delay=self.quart_config.quart_reload_delay,
reload_includes=self.quart_config.quart_reload_includes,
reload_excludes=self.quart_config.quart_reload_excludes,
debug=self.config.debug,
log_config=LOGGING_CONFIG,
**kwargs)
**kwargs,
)
async def _handle_http(self, adapter: str):
request: Request = _request
data: bytes = await request.get_data() # type: ignore
if adapter not in self._adapters:
logger.warning(f'Unknown adapter {adapter}. '
'Please register the adapter before use.')
logger.warning(
f"Unknown adapter {adapter}. " "Please register the adapter before use."
)
raise exceptions.NotFound()
BotClass = self._adapters[adapter]
http_request = HTTPRequest(request.http_version, request.scheme,
request.path, request.query_string,
dict(request.headers), request.method, data)
http_request = HTTPRequest(
request.http_version,
request.scheme,
request.path,
request.query_string,
dict(request.headers),
request.method,
data,
)
self_id, response = await BotClass.check_permission(self, http_request)
if not self_id:
raise exceptions.Unauthorized(
description=(response and response.body or b"").decode())
description=(response and response.body or b"").decode()
)
if self_id in self._clients:
logger.warning("There's already a reverse websocket connection,"
"so the event may be handled twice.")
logger.warning(
"There's already a reverse websocket connection,"
"so the event may be handled twice."
)
bot = BotClass(self_id, http_request)
asyncio.create_task(bot.handle_message(data))
return Response(response and response.body or "",
response and response.status or 200)
return Response(
response and response.body or "", response and response.status or 200
)
async def _handle_ws_reverse(self, adapter: str):
websocket: QuartWebSocket = _websocket
ws = WebSocket(websocket.http_version, websocket.scheme,
websocket.path, websocket.query_string,
dict(websocket.headers), websocket)
ws = WebSocket(
websocket.http_version,
websocket.scheme,
websocket.path,
websocket.query_string,
dict(websocket.headers),
websocket,
)
if adapter not in self._adapters:
logger.warning(
f'Unknown adapter {adapter}. Please register the adapter before use.'
f"Unknown adapter {adapter}. Please register the adapter before use."
)
raise exceptions.NotFound()
@ -242,20 +259,22 @@ class Driver(ReverseDriver):
if not self_id:
raise exceptions.Unauthorized(
description=(response and response.body or b"").decode())
description=(response and response.body or b"").decode()
)
if self_id in self._clients:
logger.opt(colors=True).warning(
"There's already a websocket connection, "
f"<y>{escape_tag(adapter.upper())} Bot {escape_tag(self_id)}</y> ignored."
)
raise exceptions.Forbidden(description='Client already exists.')
raise exceptions.Forbidden(description="Client already exists.")
bot = BotClass(self_id, ws)
await ws.accept()
logger.opt(colors=True).info(
f"WebSocket Connection from <y>{escape_tag(adapter.upper())} "
f"Bot {escape_tag(self_id)}</y> Accepted!")
f"Bot {escape_tag(self_id)}</y> Accepted!"
)
self._bot_connect(bot)
try:
@ -267,7 +286,8 @@ class Driver(ReverseDriver):
break
except Exception as e:
logger.opt(exception=e).error(
"Error when receiving data from websocket.")
"Error when receiving data from websocket."
)
break
asyncio.create_task(bot.handle_message(data.encode()))