ayaka.driver.ayakabot.driver 

  1from importlib import import_module
  2import json
  3import asyncio
  4import uvicorn
  5from loguru import logger
  6from fastapi import FastAPI, WebSocket
  7from pathlib import Path
  8import re
  9
 10from .bot import Bot
 11from .event import json_to_event, MessageEvent
 12from .websocket import FastAPIWebSocket
 13
 14app = FastAPI()
 15
 16
 17class Config:
 18    port = 19900
 19    ayaka_prefix = "#"
 20    ayaka_separate = " "
 21    ayaka_exclude_old = True
 22    fastapi_reload = True
 23
 24
 25class Driver:
 26    connect_calls = []
 27    disconnect_calls = []
 28    deal_event = None
 29    config = Config()
 30
 31    def deal(self, bot, event):
 32        if self.deal_event:
 33            if isinstance(event, MessageEvent):
 34                logger.success(str(event))
 35                asyncio.create_task(self.deal_event(bot, event))
 36
 37    async def bot_connect(self, bot: Bot):
 38        ts = []
 39        for call in self.connect_calls:
 40            ts.append(asyncio.create_task(call(bot)))
 41        await asyncio.gather(*ts)
 42
 43    async def bot_disconnect(self, bot: Bot):
 44        ts = []
 45        for call in self.disconnect_calls:
 46            ts.append(asyncio.create_task(call(bot)))
 47        await asyncio.gather(*ts)
 48
 49    def on_bot_connect(self, func):
 50        self.connect_calls.append(func)
 51
 52    def on_bot_disconnect(self, func):
 53        self.disconnect_calls.append(func)
 54
 55    def on_startup(self, func):
 56        app.on_event("startup")(func)
 57
 58    def on_shutdown(self, func):
 59        app.on_event("shutdown")(func)
 60
 61
 62driver = Driver()
 63
 64
 65# 启动服务
 66def run(host="127.0.0.1", port=driver.config.port, reload=True):
 67    uvicorn.run(
 68        app=f"{__name__}:app",
 69        host=host,
 70        port=port,
 71        reload=reload,
 72    )
 73
 74
 75def load_plugins(path):
 76    path = Path(path)
 77    for p in path.iterdir():
 78        if p.name.startswith("_"):
 79            continue
 80
 81        load_plugin(p)
 82
 83
 84def load_plugin(path):
 85    if isinstance(path, Path):
 86        p = path
 87    else:
 88        p = Path(path)
 89
 90    name = re.sub(r"\\|/", ".", str(p))
 91    try:
 92        import_module(name)
 93        logger.opt(colors=True).success(f"导入成功 \"<y>{p.stem}</y>\"")
 94    except:
 95        logger.opt(colors=True).exception(f"导入失败 \"<y>{p.stem}</y>\"")
 96
 97
 98@app.websocket("/onebot/v11/ws")
 99async def endpoint(websocket: WebSocket):
100    self_id = websocket.headers.get("x-self-id")
101    ws = FastAPIWebSocket(websocket)
102    bot = Bot(ws, self_id)
103
104    # 建立ws连接
105    await ws.accept()
106    await driver.bot_connect(bot)
107
108    try:
109        # 监听循环
110        while True:
111            data = await ws.receive()
112            json_data = json.loads(data)
113            # 将json解析为对应的event
114            event = json_to_event(json_data)
115
116            if not event:
117                continue
118
119            driver.deal(bot, event)
120
121    except:
122        logger.exception("连接中断")
123    finally:
124        # 结束ws连接
125        await driver.bot_disconnect(bot)
126        await ws.close()
127
128
129def get_driver():
130    return driver
131
132
133def on_message(priority, block, handlers):
134    driver.deal_event = handlers[0]
class Config: 
18class Config:
19    port = 19900
20    ayaka_prefix = "#"
21    ayaka_separate = " "
22    ayaka_exclude_old = True
23    fastapi_reload = True
Config()
class Driver: 
26class Driver:
27    connect_calls = []
28    disconnect_calls = []
29    deal_event = None
30    config = Config()
31
32    def deal(self, bot, event):
33        if self.deal_event:
34            if isinstance(event, MessageEvent):
35                logger.success(str(event))
36                asyncio.create_task(self.deal_event(bot, event))
37
38    async def bot_connect(self, bot: Bot):
39        ts = []
40        for call in self.connect_calls:
41            ts.append(asyncio.create_task(call(bot)))
42        await asyncio.gather(*ts)
43
44    async def bot_disconnect(self, bot: Bot):
45        ts = []
46        for call in self.disconnect_calls:
47            ts.append(asyncio.create_task(call(bot)))
48        await asyncio.gather(*ts)
49
50    def on_bot_connect(self, func):
51        self.connect_calls.append(func)
52
53    def on_bot_disconnect(self, func):
54        self.disconnect_calls.append(func)
55
56    def on_startup(self, func):
57        app.on_event("startup")(func)
58
59    def on_shutdown(self, func):
60        app.on_event("shutdown")(func)
Driver()
def deal(self, bot, event): 
32    def deal(self, bot, event):
33        if self.deal_event:
34            if isinstance(event, MessageEvent):
35                logger.success(str(event))
36                asyncio.create_task(self.deal_event(bot, event))
async def bot_connect(self, bot: ayaka.driver.ayakabot.bot.Bot): 
38    async def bot_connect(self, bot: Bot):
39        ts = []
40        for call in self.connect_calls:
41            ts.append(asyncio.create_task(call(bot)))
42        await asyncio.gather(*ts)
async def bot_disconnect(self, bot: ayaka.driver.ayakabot.bot.Bot): 
44    async def bot_disconnect(self, bot: Bot):
45        ts = []
46        for call in self.disconnect_calls:
47            ts.append(asyncio.create_task(call(bot)))
48        await asyncio.gather(*ts)
def on_bot_connect(self, func): 
50    def on_bot_connect(self, func):
51        self.connect_calls.append(func)
def on_bot_disconnect(self, func): 
53    def on_bot_disconnect(self, func):
54        self.disconnect_calls.append(func)
def on_startup(self, func): 
56    def on_startup(self, func):
57        app.on_event("startup")(func)
def on_shutdown(self, func): 
59    def on_shutdown(self, func):
60        app.on_event("shutdown")(func)
def run(host='127.0.0.1', port=19900, reload=True): 
67def run(host="127.0.0.1", port=driver.config.port, reload=True):
68    uvicorn.run(
69        app=f"{__name__}:app",
70        host=host,
71        port=port,
72        reload=reload,
73    )
def load_plugins(path): 
76def load_plugins(path):
77    path = Path(path)
78    for p in path.iterdir():
79        if p.name.startswith("_"):
80            continue
81
82        load_plugin(p)
def load_plugin(path): 
85def load_plugin(path):
86    if isinstance(path, Path):
87        p = path
88    else:
89        p = Path(path)
90
91    name = re.sub(r"\\|/", ".", str(p))
92    try:
93        import_module(name)
94        logger.opt(colors=True).success(f"导入成功 \"<y>{p.stem}</y>\"")
95    except:
96        logger.opt(colors=True).exception(f"导入失败 \"<y>{p.stem}</y>\"")
@app.websocket('/onebot/v11/ws')
async def endpoint(websocket: starlette.websockets.WebSocket): 
 99@app.websocket("/onebot/v11/ws")
100async def endpoint(websocket: WebSocket):
101    self_id = websocket.headers.get("x-self-id")
102    ws = FastAPIWebSocket(websocket)
103    bot = Bot(ws, self_id)
104
105    # 建立ws连接
106    await ws.accept()
107    await driver.bot_connect(bot)
108
109    try:
110        # 监听循环
111        while True:
112            data = await ws.receive()
113            json_data = json.loads(data)
114            # 将json解析为对应的event
115            event = json_to_event(json_data)
116
117            if not event:
118                continue
119
120            driver.deal(bot, event)
121
122    except:
123        logger.exception("连接中断")
124    finally:
125        # 结束ws连接
126        await driver.bot_disconnect(bot)
127        await ws.close()
def get_driver(): 
130def get_driver():
131    return driver
def on_message(priority, block, handlers): 
134def on_message(priority, block, handlers):
135    driver.deal_event = handlers[0]