forked from bot/app
✨ 对通道类添加类型检查和泛型
This commit is contained in:
parent
03057c8ef9
commit
f980e77a4a
@ -10,11 +10,9 @@ Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
|
|||||||
|
|
||||||
本模块定义了一个通用的通道类,用于进程间通信
|
本模块定义了一个通用的通道类,用于进程间通信
|
||||||
"""
|
"""
|
||||||
import functools
|
|
||||||
import multiprocessing
|
|
||||||
import threading
|
import threading
|
||||||
from multiprocessing import Pipe
|
from multiprocessing import Pipe
|
||||||
from typing import Any, Optional, Callable, Awaitable, List, TypeAlias
|
from typing import Any, Awaitable, Callable, Optional, TypeAlias, TypeVar, Generic, get_args
|
||||||
from uuid import uuid4
|
from uuid import uuid4
|
||||||
|
|
||||||
from liteyuki.utils import IS_MAIN_PROCESS, is_coroutine_callable, run_coroutine
|
from liteyuki.utils import IS_MAIN_PROCESS, is_coroutine_callable, run_coroutine
|
||||||
@ -30,18 +28,21 @@ FILTER_FUNC: TypeAlias = SYNC_FILTER_FUNC | ASYNC_FILTER_FUNC
|
|||||||
_channel: dict[str, "Channel"] = {}
|
_channel: dict[str, "Channel"] = {}
|
||||||
_callback_funcs: dict[str, ON_RECEIVE_FUNC] = {}
|
_callback_funcs: dict[str, ON_RECEIVE_FUNC] = {}
|
||||||
|
|
||||||
"""子进程可用的主动和被动通道"""
|
T = TypeVar("T")
|
||||||
active_channel: Optional["Channel"] = None
|
|
||||||
passive_channel: Optional["Channel"] = None
|
|
||||||
|
|
||||||
|
|
||||||
class Channel:
|
class Channel(Generic[T]):
|
||||||
"""
|
"""
|
||||||
通道类,可以在进程间和进程内通信,双向但同时只能有一个发送者和一个接收者
|
通道类,可以在进程间和进程内通信,双向但同时只能有一个发送者和一个接收者
|
||||||
有两种接收工作方式,但是只能选择一种,主动接收和被动接收,主动接收使用 `receive` 方法,被动接收使用 `on_receive` 装饰器
|
有两种接收工作方式,但是只能选择一种,主动接收和被动接收,主动接收使用 `receive` 方法,被动接收使用 `on_receive` 装饰器
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, _id: str):
|
def __init__(self, _id: str, type_check: bool = False):
|
||||||
|
"""
|
||||||
|
初始化通道
|
||||||
|
Args:
|
||||||
|
_id: 通道ID
|
||||||
|
"""
|
||||||
self.conn_send, self.conn_recv = Pipe()
|
self.conn_send, self.conn_recv = Pipe()
|
||||||
self._closed = False
|
self._closed = False
|
||||||
self._on_main_receive_funcs: list[str] = []
|
self._on_main_receive_funcs: list[str] = []
|
||||||
@ -51,20 +52,67 @@ class Channel:
|
|||||||
self.is_main_receive_loop_running = False
|
self.is_main_receive_loop_running = False
|
||||||
self.is_sub_receive_loop_running = False
|
self.is_sub_receive_loop_running = False
|
||||||
|
|
||||||
|
if type_check:
|
||||||
|
if self._get_generic_type() is None:
|
||||||
|
raise TypeError("Type hint is required for enforcing type check.")
|
||||||
|
self.type_check = type_check
|
||||||
|
|
||||||
|
def _get_generic_type(self) -> Optional[type]:
|
||||||
|
"""
|
||||||
|
获取通道传递泛型类型
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Optional[type]: 泛型类型
|
||||||
|
"""
|
||||||
|
if hasattr(self, '__orig_class__'):
|
||||||
|
return get_args(self.__orig_class__)[0]
|
||||||
|
return None
|
||||||
|
|
||||||
|
def _validate_structure(self, data: any, structure: type | tuple | list | dict) -> bool:
|
||||||
|
"""
|
||||||
|
验证数据结构
|
||||||
|
Args:
|
||||||
|
data: 数据
|
||||||
|
structure: 结构
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
bool: 是否通过验证
|
||||||
|
"""
|
||||||
|
if isinstance(structure, type):
|
||||||
|
return isinstance(data, structure)
|
||||||
|
elif isinstance(structure, tuple):
|
||||||
|
if not isinstance(data, tuple) or len(data) != len(structure):
|
||||||
|
return False
|
||||||
|
return all(self._validate_structure(d, s) for d, s in zip(data, structure))
|
||||||
|
elif isinstance(structure, list):
|
||||||
|
if not isinstance(data, list):
|
||||||
|
return False
|
||||||
|
return all(self._validate_structure(d, structure[0]) for d in data)
|
||||||
|
elif isinstance(structure, dict):
|
||||||
|
if not isinstance(data, dict):
|
||||||
|
return False
|
||||||
|
return all(k in data and self._validate_structure(data[k], structure[k]) for k in structure)
|
||||||
|
return False
|
||||||
|
|
||||||
def __str__(self):
|
def __str__(self):
|
||||||
return f"Channel({self.name})"
|
return f"Channel({self.name})"
|
||||||
|
|
||||||
def send(self, data: Any):
|
def send(self, data: T):
|
||||||
"""
|
"""
|
||||||
发送数据
|
发送数据
|
||||||
Args:
|
Args:
|
||||||
data: 数据
|
data: 数据
|
||||||
"""
|
"""
|
||||||
|
if self.type_check:
|
||||||
|
_type = self._get_generic_type()
|
||||||
|
if not self._validate_structure(data, _type):
|
||||||
|
raise TypeError(f"Data must be an instance of {_type}, {type(data)} found")
|
||||||
|
|
||||||
if self._closed:
|
if self._closed:
|
||||||
raise RuntimeError("Cannot send to a closed channel")
|
raise RuntimeError("Cannot send to a closed channel")
|
||||||
self.conn_send.send(data)
|
self.conn_send.send(data)
|
||||||
|
|
||||||
def receive(self) -> Any:
|
def receive(self) -> T:
|
||||||
"""
|
"""
|
||||||
接收数据
|
接收数据
|
||||||
Args:
|
Args:
|
||||||
@ -168,6 +216,38 @@ class Channel:
|
|||||||
return self.receive()
|
return self.receive()
|
||||||
|
|
||||||
|
|
||||||
|
"""子进程可用的主动和被动通道"""
|
||||||
|
active_channel: Optional["Channel"] = None
|
||||||
|
passive_channel: Optional["Channel"] = None
|
||||||
|
if not IS_MAIN_PROCESS:
|
||||||
|
"""sub process only"""
|
||||||
|
active_channel: Optional["Channel"] = None
|
||||||
|
passive_channel: Optional["Channel"] = None
|
||||||
|
|
||||||
|
"""通道传递通道,主进程单例,子进程初始化时实例化"""
|
||||||
|
channel_deliver_active_channel: Optional["Channel"]
|
||||||
|
channel_deliver_passive_channel: Optional["Channel"]
|
||||||
|
if IS_MAIN_PROCESS:
|
||||||
|
channel_deliver_active_channel: Optional["Channel"] = Channel(_id="channel_deliver_active_channel")
|
||||||
|
channel_deliver_passive_channel: Optional["Channel"] = Channel(_id="channel_deliver_passive_channel")
|
||||||
|
|
||||||
|
|
||||||
|
@channel_deliver_passive_channel.on_receive(filter_func=lambda data: data[0] == "set_channel")
|
||||||
|
def on_set_channel(data: tuple[str, str, Channel, Channel]):
|
||||||
|
name, channel, temp_channel = data[1:]
|
||||||
|
temp_channel.send(set_channel(name, channel))
|
||||||
|
|
||||||
|
|
||||||
|
@channel_deliver_active_channel.on_receive(filter_func=lambda data: data[0] == "get_channel")
|
||||||
|
def on_get_channel(data: tuple[str, Channel]):
|
||||||
|
name = data[1:]
|
||||||
|
channel = get_channel()
|
||||||
|
return channel
|
||||||
|
else:
|
||||||
|
channel_deliver_active_channel = None
|
||||||
|
channel_deliver_passive_channel = None
|
||||||
|
|
||||||
|
|
||||||
def set_channel(name: str, channel: Channel):
|
def set_channel(name: str, channel: Channel):
|
||||||
"""
|
"""
|
||||||
设置通道实例
|
设置通道实例
|
||||||
@ -175,12 +255,16 @@ def set_channel(name: str, channel: Channel):
|
|||||||
name: 通道名称
|
name: 通道名称
|
||||||
channel: 通道实例
|
channel: 通道实例
|
||||||
"""
|
"""
|
||||||
if not IS_MAIN_PROCESS:
|
|
||||||
raise RuntimeError(f"Function {__name__} should only be called in the main process.")
|
|
||||||
|
|
||||||
if not isinstance(channel, Channel):
|
if not isinstance(channel, Channel):
|
||||||
raise TypeError(f"channel must be an instance of Channel, {type(channel)} found")
|
raise TypeError(f"channel must be an instance of Channel, {type(channel)} found")
|
||||||
|
|
||||||
|
if IS_MAIN_PROCESS:
|
||||||
_channel[name] = channel
|
_channel[name] = channel
|
||||||
|
else:
|
||||||
|
# 请求主进程设置通道
|
||||||
|
temp_channel = Channel(_id="temp_channel")
|
||||||
|
channel_deliver_passive_channel.send(("set_channel", name, channel, temp_channel))
|
||||||
|
return temp_channel.receive()
|
||||||
|
|
||||||
|
|
||||||
def set_channels(channels: dict[str, Channel]):
|
def set_channels(channels: dict[str, Channel]):
|
||||||
|
@ -197,7 +197,6 @@ shared_memory: Optional[KeyValueStore] = None
|
|||||||
if IS_MAIN_PROCESS:
|
if IS_MAIN_PROCESS:
|
||||||
shared_memory = GlobalKeyValueStore.get_instance()
|
shared_memory = GlobalKeyValueStore.get_instance()
|
||||||
|
|
||||||
|
|
||||||
@shared_memory.passive_chan.on_receive(lambda d: d[0] == "get")
|
@shared_memory.passive_chan.on_receive(lambda d: d[0] == "get")
|
||||||
def on_get(data: tuple[str, str, any, Channel]):
|
def on_get(data: tuple[str, str, any, Channel]):
|
||||||
data[3].send(shared_memory.get(data[1], data[2]))
|
data[3].send(shared_memory.get(data[1], data[2]))
|
||||||
|
@ -16,17 +16,19 @@ import threading
|
|||||||
from multiprocessing import Process
|
from multiprocessing import Process
|
||||||
from typing import Any, Callable, TYPE_CHECKING, TypeAlias
|
from typing import Any, Callable, TYPE_CHECKING, TypeAlias
|
||||||
|
|
||||||
from liteyuki.comm import Channel, get_channel, set_channels
|
from liteyuki.comm.channel import Channel, get_channel, set_channels
|
||||||
from liteyuki.comm.storage import shared_memory
|
from liteyuki.comm.storage import shared_memory
|
||||||
from liteyuki.log import logger
|
from liteyuki.log import logger
|
||||||
from liteyuki.utils import IS_MAIN_PROCESS
|
from liteyuki.utils import IS_MAIN_PROCESS
|
||||||
|
|
||||||
TARGET_FUNC: TypeAlias = Callable[..., Any]
|
if IS_MAIN_PROCESS:
|
||||||
|
from liteyuki.comm.channel import channel_deliver_active_channel, channel_deliver_passive_channel
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from liteyuki.bot import LiteyukiBot
|
from liteyuki.bot import LiteyukiBot
|
||||||
from liteyuki.comm.storage import KeyValueStore
|
from liteyuki.comm.storage import KeyValueStore
|
||||||
|
|
||||||
|
TARGET_FUNC: TypeAlias = Callable[..., Any]
|
||||||
TIMEOUT = 10
|
TIMEOUT = 10
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
@ -34,18 +36,21 @@ __all__ = [
|
|||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
# Update the delivery_channel_wrapper function to return the top-level wrapper
|
# 函数处理一些跨进程通道的
|
||||||
def _delivery_channel_wrapper(func: TARGET_FUNC, chan_active: Channel, chan_passive: Channel, sm: "KeyValueStore", *args, **kwargs):
|
def _delivery_channel_wrapper(func: TARGET_FUNC, chan_active: Channel, chan_passive: Channel, sm: "KeyValueStore", *args, **kwargs):
|
||||||
"""
|
"""
|
||||||
子进程入口函数
|
子进程入口函数
|
||||||
|
处理一些操作
|
||||||
"""
|
"""
|
||||||
# 给子进程设置通道
|
# 给子进程设置通道
|
||||||
if IS_MAIN_PROCESS:
|
if IS_MAIN_PROCESS:
|
||||||
raise RuntimeError("Function should only be called in a sub process.")
|
raise RuntimeError("Function should only be called in a sub process.")
|
||||||
|
|
||||||
from liteyuki.comm import channel
|
from liteyuki.comm import channel # type Module
|
||||||
channel.active_channel = chan_active
|
channel.active_channel = chan_active # 子进程主动通道
|
||||||
channel.passive_channel = chan_passive
|
channel.passive_channel = chan_passive # 子进程被动通道
|
||||||
|
channel.channel_deliver_active_channel = channel_deliver_active_channel # 子进程通道传递主动通道
|
||||||
|
channel.channel_deliver_passive_channel = channel_deliver_passive_channel # 子进程通道传递被动通道
|
||||||
|
|
||||||
# 给子进程创建共享内存实例
|
# 给子进程创建共享内存实例
|
||||||
from liteyuki.comm import storage
|
from liteyuki.comm import storage
|
||||||
|
Loading…
Reference in New Issue
Block a user