From 1b692dd13ffd28ec9699c36a9807e8837a2bc15a Mon Sep 17 00:00:00 2001 From: snowy Date: Fri, 16 Aug 2024 21:43:29 +0800 Subject: [PATCH] =?UTF-8?q?:sparkles:=20=E6=96=B0=E5=A2=9E=E7=BA=BF?= =?UTF-8?q?=E7=A8=8B=E5=AE=89=E5=85=A8=E5=85=B1=E4=BA=AB=E5=86=85=E5=AD=98?= =?UTF-8?q?=E5=82=A8=E5=AD=98=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- liteyuki/bot/__init__.py | 2 +- liteyuki/comm/storage.py | 46 +++++++++++++++++-- liteyuki/core/__init__.py | 6 --- .../nonebot_launcher/__init__.py | 2 +- 4 files changed, 43 insertions(+), 13 deletions(-) diff --git a/liteyuki/bot/__init__.py b/liteyuki/bot/__init__.py index ba26122e..2c0a0f97 100644 --- a/liteyuki/bot/__init__.py +++ b/liteyuki/bot/__init__.py @@ -8,7 +8,7 @@ from typing import Any, Optional from liteyuki.bot.lifespan import (LIFESPAN_FUNC, Lifespan) from liteyuki.comm import get_channel -from liteyuki.core import IS_MAIN_PROCESS +from liteyuki.utils import IS_MAIN_PROCESS from liteyuki.core.manager import ProcessManager from liteyuki.log import init_log, logger from liteyuki.plugin import load_plugins diff --git a/liteyuki/comm/storage.py b/liteyuki/comm/storage.py index 6c815036..3453caf3 100644 --- a/liteyuki/comm/storage.py +++ b/liteyuki/comm/storage.py @@ -30,6 +30,13 @@ class KeyValueStore: self.passive_chan = Channel(_id="shared_memory-passive") def set(self, key: str, value: any) -> None: + """ + 设置键值对 + Args: + key: 键 + value: 值 + + """ if IS_MAIN_PROCESS: lock = _get_lock(key) with lock: @@ -39,6 +46,15 @@ class KeyValueStore: self.passive_chan.send(("set", key, value)) def get(self, key: str, default: Optional[any] = None) -> any: + """ + 获取键值对 + Args: + key: 键 + default: 默认值 + + Returns: + any: 值 + """ if IS_MAIN_PROCESS: lock = _get_lock(key) with lock: @@ -47,18 +63,35 @@ class KeyValueStore: self.passive_chan.send(("get", key, default)) return self.active_chan.receive() - def delete(self, key: str) -> None: + def delete(self, key: str, ignore_key_error: bool = True) -> None: + """ + 删除键值对 + Args: + key: 键 + ignore_key_error: 是否忽略键不存在的错误 + + Returns: + """ if IS_MAIN_PROCESS: lock = _get_lock(key) with lock: if key in self._store: - del self._store[key] - del _locks[key] + try: + del self._store[key] + del _locks[key] + except KeyError as e: + if not ignore_key_error: + raise e else: # 向主进程发送请求删除 self.passive_chan.send(("delete", key)) def get_all(self) -> dict[str, any]: + """ + 获取所有键值对 + Returns: + dict[str, any]: 键值对 + """ if IS_MAIN_PROCESS: return self._store else: @@ -91,9 +124,7 @@ if IS_MAIN_PROCESS: @shared_memory.passive_chan.on_receive(lambda d: d[0] == "get") def on_get(d): - print(shared_memory.get_all()) shared_memory.active_chan.send(shared_memory.get(d[1], d[2])) - print("发送数据：", shared_memory.get(d[1], d[2])) @shared_memory.passive_chan.on_receive(lambda d: d[0] == "set") @@ -104,6 +135,11 @@ if IS_MAIN_PROCESS: @shared_memory.passive_chan.on_receive(lambda d: d[0] == "delete") def on_delete(d): shared_memory.delete(d[1]) + + + @shared_memory.passive_chan.on_receive(lambda d: d[0] == "get_all") + def on_get_all(d): + shared_memory.active_chan.send(shared_memory.get_all()) else: shared_memory = None diff --git a/liteyuki/core/__init__.py b/liteyuki/core/__init__.py index 8f4d9278..0f162f6a 100644 --- a/liteyuki/core/__init__.py +++ b/liteyuki/core/__init__.py @@ -2,9 +2,3 @@ import multiprocessing from .manager import * -__all__ = [ - "IS_MAIN_PROCESS" -] - -IS_MAIN_PROCESS = multiprocessing.current_process().name == "MainProcess" - diff --git a/src/liteyuki_plugins/nonebot_launcher/__init__.py b/src/liteyuki_plugins/nonebot_launcher/__init__.py index 73fa373c..9330bd81 100644 --- a/src/liteyuki_plugins/nonebot_launcher/__init__.py +++ b/src/liteyuki_plugins/nonebot_launcher/__init__.py @@ -11,7 +11,7 @@ Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved import nonebot -from liteyuki.core import IS_MAIN_PROCESS +from liteyuki.utils import IS_MAIN_PROCESS from liteyuki.plugin import PluginMetadata from .nb_utils import adapter_manager, driver_manager