From cfce94e890be73362e40a645d395439317de666f Mon Sep 17 00:00:00 2001 From: Richard Chien Date: Mon, 16 Mar 2020 20:50:20 +0800 Subject: [PATCH] add on_startup and on_websocket_connect --- nonebot/__init__.py | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/nonebot/__init__.py b/nonebot/__init__.py index d05eb14e..5e6ae3fa 100644 --- a/nonebot/__init__.py +++ b/nonebot/__init__.py @@ -1,6 +1,6 @@ import asyncio import logging -from typing import Any, Optional +from typing import Any, Optional, Callable, Awaitable import aiocqhttp from aiocqhttp import CQHttp @@ -105,6 +105,28 @@ def run(host: Optional[str] = None, port: Optional[int] = None, get_bot().run(host=host, port=port, *args, **kwargs) +def on_startup(func: Callable[[], Awaitable[None]]) \ + -> Callable[[], Awaitable[None]]: + """ + Decorator to register a function as startup callback. + """ + return get_bot().server_app.before_serving(func) + + +def on_websocket_connect(func: Callable[[], Awaitable[None]]) \ + -> Callable[[], Awaitable[None]]: + """ + Decorator to register a function as websocket connect callback. + + Only work with CQHTTP v4.14+. + """ + + async def _func(event: aiocqhttp.Event): + await func() + + return get_bot().on_meta_event('lifecycle.connect')(_func) + + from .exceptions import * from .plugin import (load_plugin, load_plugins, load_builtin_plugins, get_loaded_plugins) @@ -119,6 +141,8 @@ from .helpers import context_id __all__ = [ 'NoneBot', 'scheduler', 'init', 'get_bot', 'run', + 'on_startup', 'on_websocket_connect', + 'CQHttpError', 'load_plugin', 'load_plugins', 'load_builtin_plugins',