import abc from copy import deepcopy from typing_extensions import Self from dataclasses import field, asdict, dataclass from typing import ( Any, Dict, List, Type, Tuple, Union, Generic, TypeVar, Iterable, Optional, SupportsIndex, overload, ) from pydantic import parse_obj_as from .template import MessageTemplate TMS = TypeVar("TMS", bound="MessageSegment") TM = TypeVar("TM", bound="Message") @dataclass class MessageSegment(abc.ABC, Generic[TM]): """消息段基类""" type: str """消息段类型""" data: Dict[str, Any] = field(default_factory=dict) """消息段数据""" @classmethod @abc.abstractmethod def get_message_class(cls) -> Type[TM]: """获取消息数组类型""" raise NotImplementedError @abc.abstractmethod def __str__(self) -> str: """该消息段所代表的 str,在命令匹配部分使用""" raise NotImplementedError def __len__(self) -> int: return len(str(self)) def __ne__(self, other: Self) -> bool: return not self == other def __add__(self: TMS, other: Union[str, TMS, Iterable[TMS]]) -> TM: return self.get_message_class()(self) + other def __radd__(self: TMS, other: Union[str, TMS, Iterable[TMS]]) -> TM: return self.get_message_class()(other) + self @classmethod def __get_validators__(cls): yield cls._validate @classmethod def _validate(cls, value) -> Self: if isinstance(value, cls): return value if not isinstance(value, dict): raise ValueError(f"Expected dict for MessageSegment, got {type(value)}") if "type" not in value: raise ValueError( f"Expected dict with 'type' for MessageSegment, got {value}" ) return cls(type=value["type"], data=value.get("data", {})) def get(self, key: str, default: Any = None): return asdict(self).get(key, default) def keys(self): return asdict(self).keys() def values(self): return asdict(self).values() def items(self): return asdict(self).items() def join(self: TMS, iterable: Iterable[Union[TMS, TM]]) -> TM: return self.get_message_class()(self).join(iterable) def copy(self) -> Self: return deepcopy(self) @abc.abstractmethod def is_text(self) -> bool: """当前消息段是否为纯文本""" raise NotImplementedError class Message(List[TMS], abc.ABC): """消息序列 参数: message: 消息内容 """ def __init__( self, message: Union[str, None, Iterable[TMS], TMS] = None, ): super().__init__() if message is None: return elif isinstance(message, str): self.extend(self._construct(message)) elif isinstance(message, MessageSegment): self.append(message) elif isinstance(message, Iterable): self.extend(message) else: self.extend(self._construct(message)) # pragma: no cover @classmethod def template(cls, format_string: Union[str, TM]) -> MessageTemplate[Self]: """创建消息模板。 用法和 `str.format` 大致相同,支持以 `Message` 对象作为消息模板并输出消息对象。 并且提供了拓展的格式化控制符, 可以通过该消息类型的 `MessageSegment` 工厂方法创建消息。 参数: format_string: 格式化模板 返回: 消息格式化器 """ return MessageTemplate(format_string, cls) @classmethod @abc.abstractmethod def get_segment_class(cls) -> Type[TMS]: """获取消息段类型""" raise NotImplementedError def __str__(self) -> str: return "".join(str(seg) for seg in self) @classmethod def __get_validators__(cls): yield cls._validate @classmethod def _validate(cls, value) -> Self: if isinstance(value, cls): return value elif isinstance(value, Message): raise ValueError(f"Type {type(value)} can not be converted to {cls}") elif isinstance(value, str): pass elif isinstance(value, dict): value = parse_obj_as(cls.get_segment_class(), value) elif isinstance(value, Iterable): value = [parse_obj_as(cls.get_segment_class(), v) for v in value] else: raise ValueError( f"Expected str, dict or iterable for Message, got {type(value)}" ) return cls(value) @staticmethod @abc.abstractmethod def _construct(msg: str) -> Iterable[TMS]: """构造消息数组""" raise NotImplementedError def __add__(self, other: Union[str, TMS, Iterable[TMS]]) -> Self: result = self.copy() result += other return result def __radd__(self, other: Union[str, TMS, Iterable[TMS]]) -> Self: result = self.__class__(other) return result + self def __iadd__(self, other: Union[str, TMS, Iterable[TMS]]) -> Self: if isinstance(other, str): self.extend(self._construct(other)) elif isinstance(other, MessageSegment): self.append(other) elif isinstance(other, Iterable): self.extend(other) else: raise TypeError(f"Unsupported type {type(other)!r}") return self @overload def __getitem__(self, args: str) -> Self: """获取仅包含指定消息段类型的消息 参数: args: 消息段类型 返回: 所有类型为 `args` 的消息段 """ @overload def __getitem__(self, args: Tuple[str, int]) -> TMS: """索引指定类型的消息段 参数: args: 消息段类型和索引 返回: 类型为 `args[0]` 的消息段第 `args[1]` 个 """ @overload def __getitem__(self, args: Tuple[str, slice]) -> Self: """切片指定类型的消息段 参数: args: 消息段类型和切片 返回: 类型为 `args[0]` 的消息段切片 `args[1]` """ @overload def __getitem__(self, args: int) -> TMS: """索引消息段 参数: args: 索引 返回: 第 `args` 个消息段 """ @overload def __getitem__(self, args: slice) -> Self: """切片消息段 参数: args: 切片 返回: 消息切片 `args` """ def __getitem__( self, args: Union[ str, Tuple[str, int], Tuple[str, slice], int, slice, ], ) -> Union[TMS, Self]: arg1, arg2 = args if isinstance(args, tuple) else (args, None) if isinstance(arg1, int) and arg2 is None: return super().__getitem__(arg1) elif isinstance(arg1, slice) and arg2 is None: return self.__class__(super().__getitem__(arg1)) elif isinstance(arg1, str) and arg2 is None: return self.__class__(seg for seg in self if seg.type == arg1) elif isinstance(arg1, str) and isinstance(arg2, int): return [seg for seg in self if seg.type == arg1][arg2] elif isinstance(arg1, str) and isinstance(arg2, slice): return self.__class__([seg for seg in self if seg.type == arg1][arg2]) else: raise ValueError("Incorrect arguments to slice") # pragma: no cover def __contains__(self, value: Union[TMS, str]) -> bool: """检查消息段是否存在 参数: value: 消息段或消息段类型 返回: 消息内是否存在给定消息段或给定类型的消息段 """ if isinstance(value, str): return bool(next((seg for seg in self if seg.type == value), None)) return super().__contains__(value) def has(self, value: Union[TMS, str]) -> bool: """与 {ref}``__contains__` ` 相同""" return value in self def index(self, value: Union[TMS, str], *args: SupportsIndex) -> int: """索引消息段 参数: value: 消息段或者消息段类型 arg: start 与 end 返回: 索引 index 异常: ValueError: 消息段不存在 """ if isinstance(value, str): first_segment = next((seg for seg in self if seg.type == value), None) if first_segment is None: raise ValueError(f"Segment with type {value!r} is not in message") return super().index(first_segment, *args) return super().index(value, *args) def get(self, type_: str, count: Optional[int] = None) -> Self: """获取指定类型的消息段 参数: type_: 消息段类型 count: 获取个数 返回: 构建的新消息 """ if count is None: return self[type_] iterator, filtered = ( seg for seg in self if seg.type == type_ ), self.__class__() for _ in range(count): seg = next(iterator, None) if seg is None: break filtered.append(seg) return filtered def count(self, value: Union[TMS, str]) -> int: """计算指定消息段的个数 参数: value: 消息段或消息段类型 返回: 个数 """ return len(self[value]) if isinstance(value, str) else super().count(value) def only(self, value: Union[TMS, str]) -> bool: """检查消息中是否仅包含指定消息段 参数: value: 指定消息段或消息段类型 返回: 是否仅包含指定消息段 """ if isinstance(value, str): return all(seg.type == value for seg in self) return all(seg == value for seg in self) def append(self, obj: Union[str, TMS]) -> Self: """添加一个消息段到消息数组末尾。 参数: obj: 要添加的消息段 """ if isinstance(obj, MessageSegment): super().append(obj) elif isinstance(obj, str): self.extend(self._construct(obj)) else: raise ValueError(f"Unexpected type: {type(obj)} {obj}") # pragma: no cover return self def extend(self, obj: Union[Self, Iterable[TMS]]) -> Self: """拼接一个消息数组或多个消息段到消息数组末尾。 参数: obj: 要添加的消息数组 """ for segment in obj: self.append(segment) return self def join(self, iterable: Iterable[Union[TMS, Self]]) -> Self: """将多个消息连接并将自身作为分割 参数: iterable: 要连接的消息 返回: 连接后的消息 """ ret = self.__class__() for index, msg in enumerate(iterable): if index != 0: ret.extend(self) if isinstance(msg, MessageSegment): ret.append(msg.copy()) else: ret.extend(msg.copy()) return ret def copy(self) -> Self: """深拷贝消息""" return deepcopy(self) def include(self, *types: str) -> Self: """过滤消息 参数: types: 包含的消息段类型 返回: 新构造的消息 """ return self.__class__(seg for seg in self if seg.type in types) def exclude(self, *types: str) -> Self: """过滤消息 参数: types: 不包含的消息段类型 返回: 新构造的消息 """ return self.__class__(seg for seg in self if seg.type not in types) def extract_plain_text(self) -> str: """提取消息内纯文本消息""" return "".join(str(seg) for seg in self if seg.is_text())