ProjectHentai/nonebot2
6
0
Fork 0
mirror of https://github.com/nonebot/nonebot2.git synced 2025-09-06 20:16:47 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

✨ Feature: 细化 driver 职责类型 (#2296)

Browse Source
This commit is contained in:
Ju4tCode
2023-08-26 11:03:24 +08:00
committed by GitHub
parent 807a86371d
commit 2e635370bb
20 changed files with 632 additions and 284 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
nonebot/__init__.py
View File

@ -53,7 +53,7 @@ from nonebot.config import Env, Config
from nonebot.log import logger as logger
from nonebot.adapters import Bot, Adapter
from nonebot.utils import escape_tag, resolve_dot_notation
from nonebot.drivers import Driver, ReverseDriver, combine_driver
from nonebot.drivers import Driver, ASGIMixin, combine_driver
try:
__version__ = version("nonebot2")
@ -149,13 +149,13 @@ def get_adapters() -> Dict[str, Adapter]:
def get_app() -> Any:
"""获取全局 {ref}`nonebot.drivers.ReverseDriver` 对应的 Server App 对象。
"""获取全局 {ref}`nonebot.drivers.ASGIMixin` 对应的 Server App 对象。
返回:
Server App 对象
异常:
AssertionError: 全局 Driver 对象不是 {ref}`nonebot.drivers.ReverseDriver` 类型
AssertionError: 全局 Driver 对象不是 {ref}`nonebot.drivers.ASGIMixin` 类型
ValueError: 全局 {ref}`nonebot.drivers.Driver` 对象尚未初始化
({ref}`nonebot.init <nonebot.init>` 尚未调用)
@ -165,21 +165,19 @@ def get_app() -> Any:
```
"""
driver = get_driver()
assert isinstance(
driver, ReverseDriver
), "app object is only available for reverse driver"
assert isinstance(driver, ASGIMixin), "app object is only available for asgi driver"
return driver.server_app
def get_asgi() -> Any:
"""获取全局 {ref}`nonebot.drivers.ReverseDriver` 对应
"""获取全局 {ref}`nonebot.drivers.ASGIMixin` 对应
[ASGI](https://asgi.readthedocs.io/) 对象。
返回:
ASGI 对象
异常:
AssertionError: 全局 Driver 对象不是 {ref}`nonebot.drivers.ReverseDriver` 类型
AssertionError: 全局 Driver 对象不是 {ref}`nonebot.drivers.ASGIMixin` 类型
ValueError: 全局 {ref}`nonebot.drivers.Driver` 对象尚未初始化
({ref}`nonebot.init <nonebot.init>` 尚未调用)
@ -190,7 +188,7 @@ def get_asgi() -> Any:
"""
driver = get_driver()
assert isinstance(
driver, ReverseDriver
driver, ASGIMixin
), "asgi object is only available for reverse driver"
return driver.asgi

12
nonebot/drivers/__init__.py
View File

@ -8,30 +8,40 @@ FrontMatter:
"""
from nonebot.internal.driver import URL as URL
from nonebot.internal.driver import Mixin as Mixin
from nonebot.internal.driver import Driver as Driver
from nonebot.internal.driver import Cookies as Cookies
from nonebot.internal.driver import Request as Request
from nonebot.internal.driver import Response as Response
from nonebot.internal.driver import ASGIMixin as ASGIMixin
from nonebot.internal.driver import WebSocket as WebSocket
from nonebot.internal.driver import HTTPVersion as HTTPVersion
from nonebot.internal.driver import ForwardMixin as ForwardMixin
from nonebot.internal.driver import ReverseMixin as ReverseMixin
from nonebot.internal.driver import ForwardDriver as ForwardDriver
from nonebot.internal.driver import ReverseDriver as ReverseDriver
from nonebot.internal.driver import combine_driver as combine_driver
from nonebot.internal.driver import HTTPClientMixin as HTTPClientMixin
from nonebot.internal.driver import HTTPServerSetup as HTTPServerSetup
from nonebot.internal.driver import WebSocketClientMixin as WebSocketClientMixin
from nonebot.internal.driver import WebSocketServerSetup as WebSocketServerSetup
__autodoc__ = {
"URL": True,
"Driver": True,
"Cookies": True,
"Request": True,
"Response": True,
"WebSocket": True,
"HTTPVersion": True,
"Driver": True,
"Mixin": True,
"ForwardMixin": True,
"ForwardDriver": True,
"HTTPClientMixin": True,
"WebSocketClientMixin": True,
"ReverseMixin": True,
"ReverseDriver": True,
"ASGIMixin": True,
"combine_driver": True,
"HTTPServerSetup": True,
"WebSocketServerSetup": True,

21
nonebot/drivers/aiohttp.py
View File

@ -16,14 +16,19 @@ FrontMatter:
"""
from typing_extensions import override
from typing import Type, AsyncGenerator
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, AsyncGenerator
from nonebot.drivers import Request, Response
from nonebot.exception import WebSocketClosed
from nonebot.drivers.none import Driver as NoneDriver
from nonebot.drivers import WebSocket as BaseWebSocket
from nonebot.drivers import HTTPVersion, ForwardMixin, ForwardDriver, combine_driver
from nonebot.drivers import (
HTTPVersion,
HTTPClientMixin,
WebSocketClientMixin,
combine_driver,
)
try:
import aiohttp
@ -34,7 +39,7 @@ except ModuleNotFoundError as e: # pragma: no cover
) from e
class Mixin(ForwardMixin):
class Mixin(HTTPClientMixin, WebSocketClientMixin):
"""AIOHTTP Mixin"""
@property
@ -172,5 +177,11 @@ class WebSocket(BaseWebSocket):
await self.websocket.send_bytes(data)
Driver: Type[ForwardDriver] = combine_driver(NoneDriver, Mixin) # type: ignore
"""AIOHTTP Driver"""
if TYPE_CHECKING:
class Driver(Mixin, NoneDriver):
...
else:
Driver = combine_driver(NoneDriver, Mixin)
"""AIOHTTP Driver"""

8
nonebot/drivers/fastapi.py
View File

@ -25,12 +25,14 @@ from typing import Any, Dict, List, Tuple, Union, Optional
from pydantic import BaseSettings
from nonebot.config import Env
from nonebot.drivers import ASGIMixin
from nonebot.exception import WebSocketClosed
from nonebot.internal.driver import FileTypes
from nonebot.drivers import Driver as BaseDriver
from nonebot.config import Config as NoneBotConfig
from nonebot.drivers import Request as BaseRequest
from nonebot.drivers import WebSocket as BaseWebSocket
from nonebot.drivers import ReverseDriver, HTTPServerSetup, WebSocketServerSetup
from nonebot.drivers import HTTPServerSetup, WebSocketServerSetup
from ._lifespan import LIFESPAN_FUNC, Lifespan
@ -87,7 +89,7 @@ class Config(BaseSettings):
extra = "ignore"
class Driver(ReverseDriver):
class Driver(BaseDriver, ASGIMixin):
"""FastAPI 驱动框架。"""
def __init__(self, env: Env, config: NoneBotConfig):
@ -179,7 +181,7 @@ class Driver(ReverseDriver):
**kwargs,
):
"""使用 `uvicorn` 启动 FastAPI"""
super().run(host, port, app, **kwargs)
super().run(host, port, app=app, **kwargs)
LOGGING_CONFIG = {
"version": 1,
"disable_existing_loggers": False,

23
nonebot/drivers/httpx.py
View File

@ -15,18 +15,15 @@ FrontMatter:
description: nonebot.drivers.httpx 模块
"""
from typing import TYPE_CHECKING
from typing_extensions import override
from typing import Type, AsyncGenerator
from contextlib import asynccontextmanager
from nonebot.drivers.none import Driver as NoneDriver
from nonebot.drivers import (
Request,
Response,
WebSocket,
HTTPVersion,
ForwardMixin,
ForwardDriver,
HTTPClientMixin,
combine_driver,
)
@ -39,7 +36,7 @@ except ModuleNotFoundError as e: # pragma: no cover
) from e
class Mixin(ForwardMixin):
class Mixin(HTTPClientMixin):
"""HTTPX Mixin"""
@property
@ -72,12 +69,12 @@ class Mixin(ForwardMixin):
request=setup,
)
@override
@asynccontextmanager
async def websocket(self, setup: Request) -> AsyncGenerator[WebSocket, None]:
async with super().websocket(setup) as ws:
yield ws
if TYPE_CHECKING:
Driver: Type[ForwardDriver] = combine_driver(NoneDriver, Mixin) # type: ignore
"""HTTPX Driver"""
class Driver(Mixin, NoneDriver):
...
else:
Driver = combine_driver(NoneDriver, Mixin)
"""HTTPX Driver"""

6
nonebot/drivers/quart.py
View File

@ -34,12 +34,14 @@ from typing import (
from pydantic import BaseSettings
from nonebot.config import Env
from nonebot.drivers import ASGIMixin
from nonebot.exception import WebSocketClosed
from nonebot.internal.driver import FileTypes
from nonebot.drivers import Driver as BaseDriver
from nonebot.config import Config as NoneBotConfig
from nonebot.drivers import Request as BaseRequest
from nonebot.drivers import WebSocket as BaseWebSocket
from nonebot.drivers import ReverseDriver, HTTPServerSetup, WebSocketServerSetup
from nonebot.drivers import HTTPServerSetup, WebSocketServerSetup
try:
import uvicorn
@ -89,7 +91,7 @@ class Config(BaseSettings):
extra = "ignore"
class Driver(ReverseDriver):
class Driver(BaseDriver, ASGIMixin):
"""Quart 驱动框架"""
def __init__(self, env: Env, config: NoneBotConfig):

22
nonebot/drivers/websockets.py
View File

@ -19,14 +19,14 @@ import logging
from functools import wraps
from contextlib import asynccontextmanager
from typing_extensions import ParamSpec, override
from typing import Type, Union, TypeVar, Callable, Awaitable, AsyncGenerator
from typing import TYPE_CHECKING, Union, TypeVar, Callable, Awaitable, AsyncGenerator
from nonebot.drivers import Request
from nonebot.log import LoguruHandler
from nonebot.drivers import Request, Response
from nonebot.exception import WebSocketClosed
from nonebot.drivers.none import Driver as NoneDriver
from nonebot.drivers import WebSocket as BaseWebSocket
from nonebot.drivers import ForwardMixin, ForwardDriver, combine_driver
from nonebot.drivers import WebSocketClientMixin, combine_driver
try:
from websockets.exceptions import ConnectionClosed
@ -58,7 +58,7 @@ def catch_closed(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
return decorator
class Mixin(ForwardMixin):
class Mixin(WebSocketClientMixin):
"""Websockets Mixin"""
@property
@ -66,10 +66,6 @@ class Mixin(ForwardMixin):
def type(self) -> str:
return "websockets"
@override
async def request(self, setup: Request) -> Response:
return await super().request(setup)
@override
@asynccontextmanager
async def websocket(self, setup: Request) -> AsyncGenerator["WebSocket", None]:
@ -133,5 +129,11 @@ class WebSocket(BaseWebSocket):
await self.websocket.send(data)
Driver: Type[ForwardDriver] = combine_driver(NoneDriver, Mixin) # type: ignore
"""Websockets Driver"""
if TYPE_CHECKING:
class Driver(Mixin, NoneDriver):
...
else:
Driver = combine_driver(NoneDriver, Mixin)
"""Websockets Driver"""

13
nonebot/internal/adapter/adapter.py
View File

@ -7,10 +7,11 @@ from nonebot.internal.driver import (
Driver,
Request,
Response,
ASGIMixin,
WebSocket,
ForwardDriver,
ReverseDriver,
HTTPClientMixin,
HTTPServerSetup,
WebSocketClientMixin,
WebSocketServerSetup,
)
@ -72,26 +73,26 @@ class Adapter(abc.ABC):
def setup_http_server(self, setup: HTTPServerSetup):
"""设置一个 HTTP 服务器路由配置"""
if not isinstance(self.driver, ReverseDriver):
if not isinstance(self.driver, ASGIMixin):
raise TypeError("Current driver does not support http server")
self.driver.setup_http_server(setup)
def setup_websocket_server(self, setup: WebSocketServerSetup):
"""设置一个 WebSocket 服务器路由配置"""
if not isinstance(self.driver, ReverseDriver):
if not isinstance(self.driver, ASGIMixin):
raise TypeError("Current driver does not support websocket server")
self.driver.setup_websocket_server(setup)
async def request(self, setup: Request) -> Response:
"""进行一个 HTTP 客户端请求"""
if not isinstance(self.driver, ForwardDriver):
if not isinstance(self.driver, HTTPClientMixin):
raise TypeError("Current driver does not support http client")
return await self.driver.request(setup)
@asynccontextmanager
async def websocket(self, setup: Request) -> AsyncGenerator[WebSocket, None]:
"""建立一个 WebSocket 客户端连接请求"""
if not isinstance(self.driver, ForwardDriver):
if not isinstance(self.driver, WebSocketClientMixin):
raise TypeError("Current driver does not support websocket client")
async with self.driver.websocket(setup) as ws:
yield ws

5
nonebot/internal/driver/__init__.py
View File

@ -1,4 +1,5 @@
from .model import URL as URL
from .driver import Mixin as Mixin
from .model import RawURL as RawURL
from .driver import Driver as Driver
from .model import Cookies as Cookies
@ -8,6 +9,7 @@ from .model import Response as Response
from .model import DataTypes as DataTypes
from .model import FileTypes as FileTypes
from .model import WebSocket as WebSocket
from .driver import ASGIMixin as ASGIMixin
from .model import FilesTypes as FilesTypes
from .model import QueryTypes as QueryTypes
from .model import CookieTypes as CookieTypes
@ -17,9 +19,12 @@ from .model import HeaderTypes as HeaderTypes
from .model import SimpleQuery as SimpleQuery
from .model import ContentTypes as ContentTypes
from .driver import ForwardMixin as ForwardMixin
from .driver import ReverseMixin as ReverseMixin
from .model import QueryVariable as QueryVariable
from .driver import ForwardDriver as ForwardDriver
from .driver import ReverseDriver as ReverseDriver
from .driver import combine_driver as combine_driver
from .model import HTTPServerSetup as HTTPServerSetup
from .driver import HTTPClientMixin as HTTPClientMixin
from .model import WebSocketServerSetup as WebSocketServerSetup
from .driver import WebSocketClientMixin as WebSocketClientMixin

111
nonebot/internal/driver/driver.py
View File

@ -1,7 +1,19 @@
import abc
import asyncio
from typing_extensions import TypeAlias
from contextlib import AsyncExitStack, asynccontextmanager
from typing import TYPE_CHECKING, Any, Set, Dict, Type, Callable, AsyncGenerator
from typing import (
TYPE_CHECKING,
Any,
Set,
Dict,
Type,
Union,
TypeVar,
Callable,
AsyncGenerator,
overload,
)
from nonebot.log import logger
from nonebot.config import Env, Config
@ -21,11 +33,15 @@ if TYPE_CHECKING:
from nonebot.internal.adapter import Bot, Adapter
D = TypeVar("D", bound="Driver")
BOT_HOOK_PARAMS = [DependParam, BotParam, DefaultParam]
class Driver(abc.ABC):
"""Driver 基类。
"""驱动器基类。
驱动器控制框架的启动和停止，适配器的注册，以及机器人生命周期管理。
参数:
env: 包含环境信息的 Env 对象
@ -45,6 +61,7 @@ class Driver(abc.ABC):
self.config: Config = config
"""全局配置对象"""
self._bots: Dict[str, "Bot"] = {}
self._bot_tasks: Set[asyncio.Task] = set()
def __repr__(self) -> str:
return (
@ -94,6 +111,8 @@ class Driver(abc.ABC):
f"<g>Loaded adapters: {escape_tag(', '.join(self._adapters))}</g>"
)
self.on_shutdown(self._cleanup)
@abc.abstractmethod
def on_startup(self, func: Callable) -> Callable:
"""注册一个在驱动器启动时执行的函数"""
@ -156,7 +175,9 @@ class Driver(abc.ABC):
"</bg #f8bbd0></r>"
)
asyncio.create_task(_run_hook(bot))
task = asyncio.create_task(_run_hook(bot))
task.add_done_callback(self._bot_tasks.discard)
self._bot_tasks.add(task)
def _bot_disconnect(self, bot: "Bot") -> None:
"""在连接断开后，调用该函数来注销 bot 对象"""
@ -183,23 +204,49 @@ class Driver(abc.ABC):
"</bg #f8bbd0></r>"
)
asyncio.create_task(_run_hook(bot))
task = asyncio.create_task(_run_hook(bot))
task.add_done_callback(self._bot_tasks.discard)
self._bot_tasks.add(task)
async def _cleanup(self) -> None:
"""清理驱动器资源"""
if self._bot_tasks:
logger.opt(colors=True).debug(
"<y>Waiting for running bot connection hooks...</y>"
)
await asyncio.gather(*self._bot_tasks, return_exceptions=True)
class ForwardMixin(abc.ABC):
"""客户端混入基类。"""
class Mixin(abc.ABC):
"""可与其他驱动器共用的混入基类。"""
@property
@abc.abstractmethod
def type(self) -> str:
"""客户端驱动类型名称"""
"""混入驱动类型名称"""
raise NotImplementedError
class ForwardMixin(Mixin):
"""客户端混入基类。"""
class ReverseMixin(Mixin):
"""服务端混入基类。"""
class HTTPClientMixin(ForwardMixin):
"""HTTP 客户端混入基类。"""
@abc.abstractmethod
async def request(self, setup: Request) -> Response:
"""发送一个 HTTP 请求"""
raise NotImplementedError
class WebSocketClientMixin(ForwardMixin):
"""WebSocket 客户端混入基类。"""
@abc.abstractmethod
@asynccontextmanager
async def websocket(self, setup: Request) -> AsyncGenerator[WebSocket, None]:
@ -208,12 +255,11 @@ class ForwardMixin(abc.ABC):
yield # used for static type checking's generator detection
class ForwardDriver(Driver, ForwardMixin):
"""客户端基类。将客户端框架封装，以满足适配器使用。"""
class ASGIMixin(ReverseMixin):
"""ASGI 服务端基类。
class ReverseDriver(Driver):
"""服务端基类。将后端框架封装，以满足适配器使用。"""
将后端框架封装，以满足适配器使用。
"""
@property
@abc.abstractmethod
@ -238,18 +284,49 @@ class ReverseDriver(Driver):
raise NotImplementedError
def combine_driver(driver: Type[Driver], *mixins: Type[ForwardMixin]) -> Type[Driver]:
ForwardDriver: TypeAlias = ForwardMixin
"""支持客户端请求的驱动器。
**Deprecated**，请使用 {ref}`nonebot.drivers.ForwardMixin` 或其子类代替。
"""
ReverseDriver: TypeAlias = ReverseMixin
"""支持服务端请求的驱动器。
**Deprecated**，请使用 {ref}`nonebot.drivers.ReverseMixin` 或其子类代替。
"""
if TYPE_CHECKING:
class CombinedDriver(Driver, Mixin):
...
@overload
def combine_driver(driver: Type[D]) -> Type[D]:
...
@overload
def combine_driver(driver: Type[D], *mixins: Type[Mixin]) -> Type["CombinedDriver"]:
...
def combine_driver(
driver: Type[D], *mixins: Type[Mixin]
) -> Union[Type[D], Type["CombinedDriver"]]:
"""将一个驱动器和多个混入类合并。"""
# check first
assert issubclass(driver, Driver), "`driver` must be subclass of Driver"
assert all(
issubclass(m, ForwardMixin) for m in mixins
), "`mixins` must be subclass of ForwardMixin"
issubclass(m, Mixin) for m in mixins
), "`mixins` must be subclass of Mixin"
if not mixins:
return driver
def type_(self: ForwardDriver) -> str:
def type_(self: "CombinedDriver") -> str:
return (
driver.type.__get__(self)
+ "+"
@ -257,5 +334,5 @@ def combine_driver(driver: Type[Driver], *mixins: Type[ForwardMixin]) -> Type[Dr
)
return type(
"CombinedDriver", (*mixins, driver, ForwardDriver), {"type": property(type_)}
"CombinedDriver", (*mixins, driver), {"type": property(type_)}
) # type: ignore