mirror of
https://github.com/nonebot/nonebot2.git
synced 2024-09-21 05:12:34 +00:00
commit
ce0e230887
@ -6,11 +6,14 @@ from typing import (
|
|||||||
Dict,
|
Dict,
|
||||||
List,
|
List,
|
||||||
Type,
|
Type,
|
||||||
|
Tuple,
|
||||||
Union,
|
Union,
|
||||||
Generic,
|
Generic,
|
||||||
Mapping,
|
Mapping,
|
||||||
TypeVar,
|
TypeVar,
|
||||||
Iterable,
|
Iterable,
|
||||||
|
Optional,
|
||||||
|
overload,
|
||||||
)
|
)
|
||||||
|
|
||||||
from ._template import MessageTemplate
|
from ._template import MessageTemplate
|
||||||
@ -180,6 +183,71 @@ class Message(List[TMS], abc.ABC):
|
|||||||
self.extend(self._construct(other))
|
self.extend(self._construct(other))
|
||||||
return self
|
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:
|
def append(self: TM, obj: Union[str, TMS]) -> TM:
|
||||||
"""
|
"""
|
||||||
添加一个消息段到消息数组末尾
|
添加一个消息段到消息数组末尾
|
||||||
|
54
tests/test_adapters/test_message.py
Normal file
54
tests/test_adapters/test_message.py
Normal file
@ -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
|
@ -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]"
|
|
@ -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
|
from pydantic import create_model
|
||||||
|
|
||||||
@ -34,8 +34,13 @@ def make_fake_message() -> Type["Message"]:
|
|||||||
return FakeMessageSegment
|
return FakeMessageSegment
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _construct(msg: str):
|
def _construct(msg: Union[str, Iterable[Mapping]]):
|
||||||
yield FakeMessageSegment.text(msg)
|
if isinstance(msg, str):
|
||||||
|
yield FakeMessageSegment.text(msg)
|
||||||
|
else:
|
||||||
|
for seg in msg:
|
||||||
|
yield FakeMessageSegment(**seg)
|
||||||
|
return
|
||||||
|
|
||||||
return FakeMessage
|
return FakeMessage
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user