🔀 Merge pull request #705

Feature: Advanced message slice support

nonebot/adapters/_message.py

@@ -6,11 +6,14 @@ from typing import (
     Dict,
     List,
     Type,
+    Tuple,
     Union,
     Generic,
     Mapping,
     TypeVar,
     Iterable,
+    Optional,
+    overload,
 )
 
 from ._template import MessageTemplate
@@ -180,6 +183,71 @@ class Message(List[TMS], abc.ABC):
             self.extend(self._construct(other))
         return self
 
+    @overload
+    def __getitem__(self: TM, __args: str) -> TM:
+        ...
+
+    @overload
+    def __getitem__(self, __args: Tuple[str, int]) -> TMS:
+        ...
+
+    @overload
+    def __getitem__(self: TM, __args: Tuple[str, slice]) -> TM:
+        ...
+
+    @overload
+    def __getitem__(self, __args: int) -> TMS:
+        ...
+
+    @overload
+    def __getitem__(self: TM, __args: slice) -> TM:
+        ...
+
+    def __getitem__(
+        self: TM,
+        args: Union[
+            str,
+            Tuple[str, int],
+            Tuple[str, slice],
+            int,
+            slice,
+        ],
+    ) -> Union[TMS, TM]:
+        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")
+
+    def index(self, value: Union[TMS, str], *args) -> int:
+        if isinstance(value, str):
+            first_segment = next((seg for seg in self if seg.type == value), None)  # type: ignore
+            return super().index(first_segment, *args)  # type: ignore
+        return super().index(value, *args)
+
+    def get(self: TM, type_: str, count: Optional[int] = None) -> TM:
+        if count is None:
+            return self[type_]
+
+        iterator, filtered = (seg for seg in self if seg.type == type_), []
+        for _ in range(count):
+            seg = next(iterator, None)
+            if seg is None:
+                break
+            filtered.append(seg)
+        return self.__class__(filtered)
+
+    def count(self, value: Union[TMS, str]) -> int:
+        return len(self[value]) if isinstance(value, str) else super().count(value)
+
     def append(self: TM, obj: Union[str, TMS]) -> TM:
         """
         添加一个消息段到消息数组末尾

tests/test_adapters/test_message.py

@@ -0,0 +1,54 @@
+from utils import make_fake_message
+
+
+def test_message_template():
+    from nonebot.adapters import MessageTemplate
+
+    Message = make_fake_message()
+
+    template = MessageTemplate("{a:custom}{b:text}{c:image}", Message)
+
+    @template.add_format_spec
+    def custom(input: str) -> str:
+        return input + "-custom!"
+
+    formatted = template.format(a="test", b="test", c="https://example.com/test")
+    assert formatted.extract_plain_text() == "test-custom!test"
+    assert str(formatted) == "test-custom!test[fake:image]"
+
+
+def test_message_slice():
+
+    Message = make_fake_message()
+    MessageSegment = Message.get_segment_class()
+
+    message = Message(
+        [
+            MessageSegment.text("test"),
+            MessageSegment.image("test2"),
+            MessageSegment.image("test3"),
+            MessageSegment.text("test4"),
+        ]
+    )
+
+    assert message[0] == MessageSegment.text("test")
+
+    assert message[0:2] == Message(
+        [MessageSegment.text("test"), MessageSegment.image("test2")]
+    )
+
+    assert message["image"] == Message(
+        [MessageSegment.image("test2"), MessageSegment.image("test3")]
+    )
+
+    assert message["image", 0] == MessageSegment.image("test2")
+    assert message["image", 0:2] == message["image"]
+
+    assert message.index(message[0]) == 0
+    assert message.index("image") == 1
+
+    assert message.get("image") == message["image"]
+    assert message.get("image", 114514) == message["image"]
+    assert message.get("image", 1) == Message([message["image", 0]])
+
+    assert message.count("image") == 2

tests/test_adapters/test_template.py

@@ -1,17 +0,0 @@
-from utils import make_fake_message
-
-
-def test_message_template():
-    from nonebot.adapters import MessageTemplate
-
-    Message = make_fake_message()
-
-    template = MessageTemplate("{a:custom}{b:text}{c:image}", Message)
-
-    @template.add_format_spec
-    def custom(input: str) -> str:
-        return input + "-custom!"
-
-    formatted = template.format(a="test", b="test", c="https://example.com/test")
-    assert formatted.extract_plain_text() == "test-custom!test"
-    assert str(formatted) == "test-custom!test[fake:image]"

tests/utils.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Type, Optional
+from typing import TYPE_CHECKING, Type, Union, Mapping, Iterable, Optional
 
 from pydantic import create_model
 
@@ -34,8 +34,13 @@ def make_fake_message() -> Type["Message"]:
             return FakeMessageSegment
 
         @staticmethod
-        def _construct(msg: str):
-            yield FakeMessageSegment.text(msg)
+        def _construct(msg: Union[str, Iterable[Mapping]]):
+            if isinstance(msg, str):
+                yield FakeMessageSegment.text(msg)
+            else:
+                for seg in msg:
+                    yield FakeMessageSegment(**seg)
+            return
 
     return FakeMessage