diff --git a/nonebot/adapters/_template.py b/nonebot/adapters/_template.py index fad17d84..82903dc5 100644 --- a/nonebot/adapters/_template.py +++ b/nonebot/adapters/_template.py @@ -4,6 +4,9 @@ from string import Formatter from typing import ( TYPE_CHECKING, Any, + Callable, + Dict, + Optional, Set, List, Type, @@ -23,6 +26,9 @@ if TYPE_CHECKING: TM = TypeVar("TM", bound="Message") TF = TypeVar("TF", str, "Message") +FormatSpecFunc = Callable[[Any], str] +FormatSpecFunc_T = TypeVar("FormatSpecFunc_T", bound=FormatSpecFunc) + class MessageTemplate(Formatter, Generic[TF]): """消息模板格式化实现类""" @@ -50,8 +56,18 @@ class MessageTemplate(Formatter, Generic[TF]): * ``template: Union[str, Message]``: 模板 * ``factory: Union[str, Message]``: 消息构造类型,默认为 `str` """ - self.template: Union[str, TF] = template + self.template: TF = template self.factory: Type[TF] = factory + self.format_specs: Dict[str, FormatSpecFunc] = {} + + def add_format_spec( + self, spec: FormatSpecFunc_T, name: Optional[str] = None + ) -> FormatSpecFunc_T: + name = name or spec.__name__ + if name in self.format_specs: + raise ValueError(f"Format spec {name} already exists!") + self.format_specs[name] = spec + return spec def format(self, *args: Any, **kwargs: Any) -> TF: """ @@ -69,7 +85,7 @@ class MessageTemplate(Formatter, Generic[TF]): else: raise TypeError("template must be a string or instance of Message!") - return msg + return msg # type:ignore def vformat( self, format_string: str, args: Sequence[Any], kwargs: Mapping[str, Any] @@ -155,25 +171,16 @@ class MessageTemplate(Formatter, Generic[TF]): ) def format_field(self, value: Any, format_spec: str) -> Any: - if issubclass(self.factory, str): - return super().format_field(value, format_spec) - - segment_class: Type[MessageSegment] = self.factory.get_segment_class() - method = getattr(segment_class, format_spec, None) - method_type = inspect.getattr_static(segment_class, format_spec, None) + formatter: Optional[FormatSpecFunc] = self.format_specs.get(format_spec) + if (formatter is None) and (not issubclass(self.factory, str)): + segment_class: Type["MessageSegment"] = self.factory.get_segment_class() + method = getattr(segment_class, format_spec, None) + if inspect.ismethod(method): + formatter = getattr(segment_class, format_spec) return ( - ( - super().format_field(value, format_spec) - if ( - (method is None) - or ( - not isinstance(method_type, (classmethod, staticmethod)) - ) # Only Call staticmethod or classmethod - ) - else method(value) - ) - if format_spec - else value + super().format_field(value, format_spec) + if formatter is None + else formatter(value) ) def _add(self, a: Any, b: Any) -> Any: diff --git a/tests/adapters/test_template.py b/tests/adapters/test_template.py new file mode 100644 index 00000000..3dbef541 --- /dev/null +++ b/tests/adapters/test_template.py @@ -0,0 +1,17 @@ +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]" diff --git a/tests/utils.py b/tests/utils.py index d516665f..a8d0c0ed 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -18,14 +18,18 @@ def make_fake_message() -> Type["Message"]: return FakeMessage def __str__(self) -> str: - return self.data["text"] + return self.data["text"] if self.type == "text" else f"[fake:{self.type}]" @classmethod def text(cls, text: str): return cls("text", {"text": text}) + @classmethod + def image(cls, url: str): + return cls("image", {"url": url}) + def is_text(self) -> bool: - return True + return self.type == "text" class FakeMessage(Message): @classmethod