ayaka.driver.ayakabot.message

  1import re
  2from io import BytesIO
  3from pathlib import Path
  4from base64 import b64encode
  5from typing import Any, Dict, List, Type, Tuple, Union, Mapping, Iterable, Optional, cast, overload
  6from copy import deepcopy
  7from dataclasses import field, asdict, dataclass
  8from pydantic import parse_obj_as
  9
 10
 11from ayaka import logger
 12from .utils import bool_to_str, escape, unescape
 13from .template import MessageTemplate
 14
 15
 16@dataclass
 17class MessageSegment:
 18    """
 19    OneBot v11 协议 MessageSegment 适配。具体方法参考协议消息段类型或源码。
 20    """
 21
 22    type: str
 23    """消息段类型"""
 24    data: Dict[str, Any] = field(default_factory=dict)
 25    """消息段数据"""
 26
 27    def __len__(self) -> int:
 28        return len(str(self))
 29
 30    def __ne__(self, other: "MessageSegment") -> bool:
 31        return not self == other
 32
 33    def __add__(self, other: Union[str, "MessageSegment", Iterable["MessageSegment"]]) -> "Message":
 34        return self.get_message_class()(self) + other
 35
 36    def __radd__(self, other: Union[str, "MessageSegment", Iterable["MessageSegment"]]) -> "Message":
 37        return self.get_message_class()(other) + self
 38
 39    @classmethod
 40    def __get_validators__(cls):
 41        yield cls._validate
 42
 43    @classmethod
 44    def _validate(cls, value):
 45        if isinstance(value, cls):
 46            return value
 47        if not isinstance(value, dict):
 48            raise ValueError(
 49                f"Expected dict for MessageSegment, got {type(value)}")
 50        return cls(**value)
 51
 52    def get(self, key: str, default: Any = None):
 53        return asdict(self).get(key, default)
 54
 55    def keys(self):
 56        return asdict(self).keys()
 57
 58    def values(self):
 59        return asdict(self).values()
 60
 61    def items(self):
 62        return asdict(self).items()
 63
 64    def copy(self) -> "MessageSegment":
 65        return deepcopy(self)
 66
 67    @classmethod
 68    def get_message_class(cls) -> Type["Message"]:
 69        return Message
 70
 71    def __str__(self) -> str:
 72        type_ = self.type
 73        data = self.data.copy()
 74
 75        # process special types
 76        if type_ == "text":
 77            # type: ignore
 78            return escape(data.get("text", ""), escape_comma=False)
 79
 80        params = ",".join(
 81            [f"{k}={escape(str(v))}" for k, v in data.items() if v is not None]
 82        )
 83        return f"[CQ:{type_}{',' if params else ''}{params}]"
 84
 85    def __add__(self, other) -> "Message":
 86        return Message(self) + (
 87            MessageSegment.text(other) if isinstance(other, str) else other
 88        )
 89
 90    def __radd__(self, other) -> "Message":
 91        return (
 92            MessageSegment.text(other) if isinstance(
 93                other, str) else Message(other)
 94        ) + self
 95
 96    def is_text(self) -> bool:
 97        return self.type == "text"
 98
 99    @staticmethod
100    def anonymous(ignore_failure: Optional[bool] = None) -> "MessageSegment":
101        return MessageSegment("anonymous", {"ignore": bool_to_str(ignore_failure)})
102
103    @staticmethod
104    def at(user_id: Union[int, str]) -> "MessageSegment":
105        return MessageSegment("at", {"qq": str(user_id)})
106
107    @staticmethod
108    def contact(type_: str, id: int) -> "MessageSegment":
109        return MessageSegment("contact", {"type": type_, "id": str(id)})
110
111    @staticmethod
112    def contact_group(group_id: int) -> "MessageSegment":
113        return MessageSegment("contact", {"type": "group", "id": str(group_id)})
114
115    @staticmethod
116    def contact_user(user_id: int) -> "MessageSegment":
117        return MessageSegment("contact", {"type": "qq", "id": str(user_id)})
118
119    @staticmethod
120    def dice() -> "MessageSegment":
121        return MessageSegment("dice", {})
122
123    @staticmethod
124    def face(id_: int) -> "MessageSegment":
125        return MessageSegment("face", {"id": str(id_)})
126
127    @staticmethod
128    def forward(id_: str) -> "MessageSegment":
129        logger.warning("Forward Message only can be received!")
130        return MessageSegment("forward", {"id": id_})
131
132    @staticmethod
133    def image(
134        file: Union[str, bytes, BytesIO, Path],
135        type_: Optional[str] = None,
136        cache: bool = True,
137        proxy: bool = True,
138        timeout: Optional[int] = None,
139    ) -> "MessageSegment":
140        if isinstance(file, BytesIO):
141            file = file.getvalue()
142        if isinstance(file, bytes):
143            file = f"base64://{b64encode(file).decode()}"
144        elif isinstance(file, Path):
145            file = f"file:///{file.resolve()}"
146        return MessageSegment(
147            "image",
148            {
149                "file": file,
150                "type": type_,
151                "cache": bool_to_str(cache),
152                "proxy": bool_to_str(proxy),
153                "timeout": timeout,
154            },
155        )
156
157    @staticmethod
158    def json(data: str) -> "MessageSegment":
159        return MessageSegment("json", {"data": data})
160
161    @staticmethod
162    def location(
163        latitude: float,
164        longitude: float,
165        title: Optional[str] = None,
166        content: Optional[str] = None,
167    ) -> "MessageSegment":
168        return MessageSegment(
169            "location",
170            {
171                "lat": str(latitude),
172                "lon": str(longitude),
173                "title": title,
174                "content": content,
175            },
176        )
177
178    @staticmethod
179    def music(type_: str, id_: int) -> "MessageSegment":
180        return MessageSegment("music", {"type": type_, "id": id_})
181
182    @staticmethod
183    def music_custom(
184        url: str,
185        audio: str,
186        title: str,
187        content: Optional[str] = None,
188        img_url: Optional[str] = None,
189    ) -> "MessageSegment":
190        return MessageSegment(
191            "music",
192            {
193                "type": "custom",
194                "url": url,
195                "audio": audio,
196                "title": title,
197                "content": content,
198                "image": img_url,
199            },
200        )
201
202    @staticmethod
203    def node(id_: int) -> "MessageSegment":
204        return MessageSegment("node", {"id": str(id_)})
205
206    @staticmethod
207    def node_custom(
208        user_id: int, nickname: str, content: Union[str, "Message"]
209    ) -> "MessageSegment":
210        return MessageSegment(
211            "node", {"user_id": str(
212                user_id), "nickname": nickname, "content": content}
213        )
214
215    @staticmethod
216    def poke(type_: str, id_: str) -> "MessageSegment":
217        return MessageSegment("poke", {"type": type_, "id": id_})
218
219    @staticmethod
220    def record(
221        file: Union[str, bytes, BytesIO, Path],
222        magic: Optional[bool] = None,
223        cache: Optional[bool] = None,
224        proxy: Optional[bool] = None,
225        timeout: Optional[int] = None,
226    ) -> "MessageSegment":
227        if isinstance(file, BytesIO):
228            file = file.getvalue()
229        if isinstance(file, bytes):
230            file = f"base64://{b64encode(file).decode()}"
231        elif isinstance(file, Path):
232            file = f"file:///{file.resolve()}"
233        return MessageSegment(
234            "record",
235            {
236                "file": file,
237                "magic": bool_to_str(magic),
238                "cache": bool_to_str(cache),
239                "proxy": bool_to_str(proxy),
240                "timeout": timeout,
241            },
242        )
243
244    @staticmethod
245    def reply(id_: int) -> "MessageSegment":
246        return MessageSegment("reply", {"id": str(id_)})
247
248    @staticmethod
249    def rps() -> "MessageSegment":
250        return MessageSegment("rps", {})
251
252    @staticmethod
253    def shake() -> "MessageSegment":
254        return MessageSegment("shake", {})
255
256    @staticmethod
257    def share(
258        url: str = "",
259        title: str = "",
260        content: Optional[str] = None,
261        image: Optional[str] = None,
262    ) -> "MessageSegment":
263        return MessageSegment(
264            "share", {"url": url, "title": title,
265                      "content": content, "image": image}
266        )
267
268    @staticmethod
269    def text(text: str) -> "MessageSegment":
270        return MessageSegment("text", {"text": text})
271
272    @staticmethod
273    def video(
274        file: Union[str, bytes, BytesIO, Path],
275        cache: Optional[bool] = None,
276        proxy: Optional[bool] = None,
277        timeout: Optional[int] = None,
278    ) -> "MessageSegment":
279        if isinstance(file, BytesIO):
280            file = file.getvalue()
281        if isinstance(file, bytes):
282            file = f"base64://{b64encode(file).decode()}"
283        elif isinstance(file, Path):
284            file = f"file:///{file.resolve()}"
285        return MessageSegment(
286            "video",
287            {
288                "file": file,
289                "cache": bool_to_str(cache),
290                "proxy": bool_to_str(proxy),
291                "timeout": timeout,
292            },
293        )
294
295    @staticmethod
296    def xml(data: str) -> "MessageSegment":
297        return MessageSegment("xml", {"data": data})
298
299
300class Message(List[MessageSegment]):
301    """
302    OneBot v11 协议 Message 适配。
303    """
304
305    def __init__(
306        self,
307        message: Union[str, None, Iterable[MessageSegment],
308                       MessageSegment] = None,
309    ):
310        super().__init__()
311        if message is None:
312            return
313        elif isinstance(message, str):
314            self.extend(self._construct(message))
315        elif isinstance(message, MessageSegment):
316            self.append(message)
317        elif isinstance(message, Iterable):
318            self.extend(message)
319        else:
320            self.extend(self._construct(message))  # pragma: no cover
321
322    @classmethod
323    def template(cls: Type["Message"], format_string: Union[str, "Message"]) -> MessageTemplate["Message"]:
324        """创建消息模板。
325
326        用法和 `str.format` 大致相同, 但是可以输出消息对象, 并且支持以 `Message` 对象作为消息模板
327
328        并且提供了拓展的格式化控制符, 可以用适用于该消息类型的 `MessageSegment` 的工厂方法创建消息
329
330        参数:
331            format_string: 格式化模板
332
333        返回:
334            消息格式化器
335        """
336        return MessageTemplate(format_string, cls)
337
338    def __str__(self) -> str:
339        return "".join(str(seg) for seg in self)
340
341    @classmethod
342    def __get_validators__(cls):
343        yield cls._validate
344
345    @classmethod
346    def _validate(cls, value):
347        if isinstance(value, cls):
348            return value
349        elif isinstance(value, Message):
350            raise ValueError(
351                f"Type {type(value)} can not be converted to {cls}")
352        elif isinstance(value, str):
353            pass
354        elif isinstance(value, dict):
355            value = parse_obj_as(cls.get_segment_class(), value)
356        elif isinstance(value, Iterable):
357            value = [parse_obj_as(cls.get_segment_class(), v) for v in value]
358        else:
359            raise ValueError(
360                f"Expected str, dict or iterable for Message, got {type(value)}"
361            )
362        return cls(value)
363
364    def __add__(self, other: Union[str, MessageSegment, Iterable[MessageSegment]]) -> "Message":
365        result = self.copy()
366        result += other
367        return result
368
369    def __radd__(self, other: Union[str, MessageSegment, Iterable[MessageSegment]]) -> "Message":
370        result = self.__class__(other)
371        return result + self
372
373    def __iadd__(self, other: Union[str, MessageSegment, Iterable[MessageSegment]]) -> "Message":
374        if isinstance(other, str):
375            self.extend(self._construct(other))
376        elif isinstance(other, MessageSegment):
377            self.append(other)
378        elif isinstance(other, Iterable):
379            self.extend(other)
380        else:
381            raise ValueError(
382                f"Unsupported type: {type(other)}")  # pragma: no cover
383        return self
384
385    @overload
386    def __getitem__(self, __args: str) -> "Message":
387        """
388        参数:
389            __args: 消息段类型
390
391        返回:
392            所有类型为 `__args` 的消息段
393        """
394
395    @overload
396    def __getitem__(self, __args: Tuple[str, int]) -> MessageSegment:
397        """
398        参数:
399            __args: 消息段类型和索引
400
401        返回:
402            类型为 `__args[0]` 的消息段第 `__args[1]` 个
403        """
404
405    @overload
406    def __getitem__(self, __args: Tuple[str, slice]) -> "Message":
407        """
408        参数:
409            __args: 消息段类型和切片
410
411        返回:
412            类型为 `__args[0]` 的消息段切片 `__args[1]`
413        """
414
415    @overload
416    def __getitem__(self, __args: int) -> MessageSegment:
417        """
418        参数:
419            __args: 索引
420
421        返回:
422            第 `__args` 个消息段
423        """
424
425    @overload
426    def __getitem__(self, __args: slice) -> "Message":
427        """
428        参数:
429            __args: 切片
430
431        返回:
432            消息切片 `__args`
433        """
434
435    def __getitem__(
436        self,
437        args: Union[
438            str,
439            Tuple[str, int],
440            Tuple[str, slice],
441            int,
442            slice,
443        ],
444    ) -> Union[MessageSegment, "Message"]:
445        arg1, arg2 = args if isinstance(args, tuple) else (args, None)
446        if isinstance(arg1, int) and arg2 is None:
447            return super().__getitem__(arg1)
448        elif isinstance(arg1, slice) and arg2 is None:
449            return self.__class__(super().__getitem__(arg1))
450        elif isinstance(arg1, str) and arg2 is None:
451            return self.__class__(seg for seg in self if seg.type == arg1)
452        elif isinstance(arg1, str) and isinstance(arg2, int):
453            return [seg for seg in self if seg.type == arg1][arg2]
454        elif isinstance(arg1, str) and isinstance(arg2, slice):
455            return self.__class__([seg for seg in self if seg.type == arg1][arg2])
456        else:
457            raise ValueError(
458                "Incorrect arguments to slice")  # pragma: no cover
459
460    def index(self, value: Union[MessageSegment, str], *args) -> int:
461        if isinstance(value, str):
462            first_segment = next(
463                (seg for seg in self if seg.type == value), None)
464            if first_segment is None:
465                raise ValueError(
466                    f"Segment with type {value} is not in message")
467            return super().index(first_segment, *args)
468        return super().index(value, *args)
469
470    def get(self, type_: str, count: Optional[int] = None) -> "Message":
471        if count is None:
472            return self[type_]
473
474        iterator, filtered = (
475            seg for seg in self if seg.type == type_
476        ), self.__class__()
477        for _ in range(count):
478            seg = next(iterator, None)
479            if seg is None:
480                break
481            filtered.append(seg)
482        return filtered
483
484    def count(self, value: Union[MessageSegment, str]) -> int:
485        return len(self[value]) if isinstance(value, str) else super().count(value)
486
487    def append(self, obj: Union[str, MessageSegment]) -> "Message":
488        """添加一个消息段到消息数组末尾。
489
490        参数:
491            obj: 要添加的消息段
492        """
493        if isinstance(obj, MessageSegment):
494            super().append(obj)
495        elif isinstance(obj, str):
496            self.extend(self._construct(obj))
497        else:
498            raise ValueError(
499                f"Unexpected type: {type(obj)} {obj}")  # pragma: no cover
500        return self
501
502    def extend(self, obj: Union["Message", Iterable[MessageSegment]]) -> "Message":
503        """拼接一个消息数组或多个消息段到消息数组末尾。
504
505        参数:
506            obj: 要添加的消息数组
507        """
508        for segment in obj:
509            self.append(segment)
510        return self
511
512    def copy(self) -> "Message":
513        return deepcopy(self)
514
515    def extract_plain_text(self) -> str:
516        """提取消息内纯文本消息"""
517
518        return "".join(str(seg) for seg in self if seg.is_text())
519
520    @classmethod
521    def get_segment_class(cls) -> Type[MessageSegment]:
522        return MessageSegment
523
524    def __add__(self, other: Union[str, Mapping, Iterable[Mapping]]) -> "Message":
525        return super(Message, self).__add__(
526            MessageSegment.text(other) if isinstance(other, str) else other
527        )
528
529    def __radd__(self, other: Union[str, Mapping, Iterable[Mapping]]) -> "Message":
530        return super(Message, self).__radd__(
531            MessageSegment.text(other) if isinstance(other, str) else other
532        )
533
534    @staticmethod
535    def _construct(
536        msg: Union[str, Mapping, Iterable[Mapping]]
537    ) -> Iterable[MessageSegment]:
538        if isinstance(msg, Mapping):
539            msg = cast(Mapping[str, Any], msg)
540            yield MessageSegment(msg["type"], msg.get("data") or {})
541            return
542        elif isinstance(msg, Iterable) and not isinstance(msg, str):
543            for seg in msg:
544                yield MessageSegment(seg["type"], seg.get("data") or {})
545            return
546        elif isinstance(msg, str):
547
548            def _iter_message(msg: str) -> Iterable[Tuple[str, str]]:
549                text_begin = 0
550                for cqcode in re.finditer(
551                    r"\[CQ:(?P<type>[a-zA-Z0-9-_.]+)"
552                    r"(?P<params>"
553                    r"(?:,[a-zA-Z0-9-_.]+=[^,\]]+)*"
554                    r"),?\]",
555                    msg,
556                ):
557                    yield "text", msg[text_begin: cqcode.pos + cqcode.start()]
558                    text_begin = cqcode.pos + cqcode.end()
559                    yield cqcode.group("type"), cqcode.group("params").lstrip(",")
560                yield "text", msg[text_begin:]
561
562            for type_, data in _iter_message(msg):
563                if type_ == "text":
564                    if data:
565                        # only yield non-empty text segment
566                        yield MessageSegment(type_, {"text": unescape(data)})
567                else:
568                    data = {
569                        k: unescape(v)
570                        for k, v in map(
571                            lambda x: x.split("=", maxsplit=1),
572                            filter(lambda x: x, (x.lstrip()
573                                   for x in data.split(","))),
574                        )
575                    }
576                    yield MessageSegment(type_, data)
577
578    def extract_plain_text(self) -> str:
579        return "".join(seg.data["text"] for seg in self if seg.is_text())