Skip to content

Commit

Permalink
✨ 迟来七个月的连接复用
Browse files Browse the repository at this point in the history
  • Loading branch information
j1g5awi committed Nov 15, 2023
2 parents 0482129 + 8be8271 commit fb02e89
Show file tree
Hide file tree
Showing 6 changed files with 390 additions and 388 deletions.
5 changes: 4 additions & 1 deletion nonebot_plugin_all4one/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,9 +36,12 @@
async def _(bot: Bot, event: Event):
if middleware := obimpl.middlewares.get(bot.self_id, None):
for event in await middleware.to_onebot_event(event):
for queue in middleware.queues:
for queue in obimpl.queues:
if queue.full():
await queue.get()
event = deepcopy(event)
if queue.self_id_prefix:
middleware.prefix_self_id(event)
await queue.put(deepcopy(event))
if a4o_config.block_event:
raise IgnoredException("All4One has transfer it to OneBot V12")
Expand Down
2 changes: 1 addition & 1 deletion nonebot_plugin_all4one/middlewares/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

from nonebot.log import logger

from .base import Queue, Middleware
from .base import Middleware

MIDDLEWARE_MAP = {}
for file in Path(__file__).parent.rglob("*.py"):
Expand Down
73 changes: 2 additions & 71 deletions nonebot_plugin_all4one/middlewares/base.py
Original file line number Diff line number Diff line change
@@ -1,65 +1,20 @@
import asyncio
from uuid import uuid4
from datetime import datetime
from abc import ABC, abstractmethod
from asyncio import Queue as BaseQueue
from typing import (
TYPE_CHECKING,
Any,
Dict,
List,
Union,
Generic,
Literal,
TypeVar,
Optional,
)
from typing import Any, Dict, List, Union, Literal, Optional

from anyio import open_file
from nonebot.adapters import Bot, Event, Message
from nonebot.adapters.onebot.v12 import UnsupportedAction
from nonebot.adapters.onebot.v12.exception import BadParam
from nonebot.adapters.onebot.v12 import Event as OneBotEvent
from nonebot.adapters.onebot.v12.event import (
Status,
BotSelf,
BotEvent,
BotStatus,
MessageEvent,
StatusUpdateMetaEvent,
)

from ..database import get_file, upload_file

_T = TypeVar("_T", bound=OneBotEvent)
if TYPE_CHECKING:

class _Queue(BaseQueue[_T]):
pass

else:

class _Queue(Generic[_T], BaseQueue):
pass


class Queue(_Queue[_T]):
def __init__(
self,
middleware: "Middleware",
maxsize: int = 0,
self_id_prefix: bool = False,
):
super().__init__(maxsize=maxsize)
self.self_id_prefix = self_id_prefix
self.middleware = middleware

async def get(self):
event = await super().get()
if self.self_id_prefix:
self.middleware.prefix_self_id(event)
return event


def supported_action(func):
"""标记支持的动作"""
Expand All @@ -70,8 +25,6 @@ def supported_action(func):
class Middleware(ABC):
def __init__(self, bot: Bot):
self.bot = bot
self.tasks: List[asyncio.Task] = []
self.queues: List[Queue[OneBotEvent]] = []
self._supported_actions = self._get_supported_actions()

def _get_supported_actions(self) -> List[str]:
Expand Down Expand Up @@ -112,29 +65,7 @@ def get_bot_self(self) -> BotSelf:
user_id=self.self_id,
)

def new_queue(
self,
self_id_prefix: bool = False,
maxsize: int = 0,
) -> Queue[OneBotEvent]:
queue = Queue(self, maxsize=maxsize, self_id_prefix=self_id_prefix)
queue.put_nowait(
StatusUpdateMetaEvent(
id=uuid4().hex,
time=datetime.now(),
type="meta",
detail_type="status_update",
sub_type="",
status=Status(
good=True,
bots=[BotStatus(self=self.get_bot_self(), online=True)],
),
)
) # TODO beta better way
self.queues.append(queue)
return queue

def prefix_self_id(self, event: Event) -> Event:
def prefix_self_id(self, event: OneBotEvent) -> OneBotEvent:
if isinstance(event, BotEvent):
event.self.user_id = "a4o@" + event.self.user_id
if isinstance(event, StatusUpdateMetaEvent):
Expand Down
Loading

0 comments on commit fb02e89

Please sign in to comment.