一种通过继承python内置Formatter实现模板字符串格式化的实现方案

[mnixry]

目前, NoneBot对MessageSegment与Message对象的字符串格式化不够完善
很多时候需要将MessageSegment转换成str再转成Message对象
这样即带来了注入的风险, 也导致了性能的额外损耗
如果单纯使用运算符重载的+号, 则无法胜任较为复杂的字符串格式化工作

所以我在这里提供一个从Python原生的string.Formatter类继承的字符串格式化类作为概念参考
经过简单测试没有发现问题

本体代码如下:

from string import Formatter
from typing import (
    Any,
    Mapping,
    Optional,
    Protocol,
    Sequence,
    Set,
    Tuple,
    TypeVar,
    Union,
)


class AddAble(Protocol):
    def __add__(self, __s: Any) -> "AddAble":
        ...

    def __str__(self) -> str:
        ...


AddAble_T = TypeVar("AddAble_T", bound=AddAble)


class MessageFormatter(Formatter):
    def format(  # type:ignore
        self,
        __format_string: str,
        *args: Sequence[AddAble],
        **kwargs: Mapping[str, AddAble],
    ) -> AddAble:
        return super().format(__format_string, *args, **kwargs)

    def vformat(
        self, format_string: str, args: Sequence[AddAble], kwargs: Mapping[str, AddAble]
    ):
        result, arg_index, used_args = self._vformat(
            format_string, args, kwargs, set(), 2
        )
        self.check_unused_args(list(used_args), args, kwargs)
        return result

    def _vformat(
        self,
        format_string: str,
        args: Sequence[Any],
        kwargs: Mapping[str, Any],
        used_args: Set[Union[int, str]],
        recursion_depth: int,
        auto_arg_index: int = 0,
    ) -> Tuple[AddAble, int, Set[Union[int, str]]]:

        if recursion_depth < 0:
            raise ValueError("Max string recursion exceeded")

        result: Optional[AddAble] = None

        for (literal_text, field_name, format_spec, conversion) in self.parse(
            format_string
        ):

            # output the literal text
            if literal_text:
                result = literal_text if result is None else (result + literal_text)

            # if there's a field, output it
            if field_name is not None:
                # this is some markup, find the object and do
                #  the formatting

                # handle arg indexing when empty field_names are given.
                if field_name == "":
                    if auto_arg_index is False:
                        raise ValueError(
                            "cannot switch from manual field specification to "
                            "automatic field numbering"
                        )
                    field_name = str(auto_arg_index)
                    auto_arg_index += 1
                elif field_name.isdigit():
                    if auto_arg_index:
                        raise ValueError(
                            "cannot switch from manual field specification to "
                            "automatic field numbering"
                        )
                    # disable auto arg incrementing, if it gets
                    # used later on, then an exception will be raised
                    auto_arg_index = False

                # given the field_name, find the object it references
                #  and the argument it came from
                obj, arg_used = self.get_field(field_name, args, kwargs)
                used_args.add(arg_used)

                assert format_spec is not None

                # do any conversion on the resulting object
                obj = self.convert_field(obj, conversion) if conversion else obj

                # expand the format spec, if needed
                format_control, auto_arg_index, formatted_args = self._vformat(
                    format_spec,
                    args,
                    kwargs,
                    used_args.copy(),
                    recursion_depth - 1,
                    auto_arg_index=auto_arg_index,
                )
                used_args |= formatted_args

                # format the object and append to the result
                formatted_text = self.format_field(obj, str(format_control))
                result = formatted_text if result is None else (result + formatted_text)

            print(f"process {result=}, {used_args=}, {field_name=}")

        return "" if result is None else result, auto_arg_index, used_args

    def format_field(self, value: AddAble_T, format_spec: str) -> Union[AddAble_T, str]:
        return super().format_field(str(value), format_spec) if format_spec else value

这是一个mah adapter可用的demo:

from nonebot.adapters.mirai import MessageEvent, Bot
from nonebot.plugin import on_message

handler = on_message()

formatter = MessageFormatter()


@handler.handle()
async def _(bot: Bot, event: MessageEvent, state: dict):
    await bot.send(event,
                   formatter.format('{prefix}{message}{suffix}',
                                    prefix='template',
                                    message=event.message_chain,
                                    suffix='yyds'),
                   at_sender=True)

欢迎大家前来进行测试和建议反馈

[mnixry]

🎉 此实现已在 nonebot/nonebot2#510 中被添加到 NoneBot2