66 lines
1.6 KiB
Python
66 lines
1.6 KiB
Python
 
|
import importlib
|
||
|
|
from collections.abc import Awaitable
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import TYPE_CHECKING, Any, TypedDict
|
||
|
|
from typing_extensions import Protocol, Unpack
|
||
|
|
|
||
|
|
from nonebot import logger
|
||
|
|
|
||
|
|
from ..config import config
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from ..bg_provider import BgBytesData
|
||
|
|
|
||
|
|
|
||
|
|
class TemplateRendererKwargs(TypedDict):
|
||
|
|
collected: dict[str, Any]
|
||
|
|
bg: "BgBytesData"
|
||
|
|
|
||
|
|
|
||
|
|
class TemplateRenderer(Protocol):
|
||
|
|
__name__: str
|
||
|
|
|
||
|
|
def __call__(
|
||
|
|
self,
|
||
|
|
**kwargs: Unpack[TemplateRendererKwargs],
|
||
|
|
) -> Awaitable[bytes]: ...
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass()
|
||
|
|
class TemplateInfo:
|
||
|
|
renderer: TemplateRenderer
|
||
|
|
collectors: set[str] | None = None
|
||
|
|
|
||
|
|
|
||
|
|
loaded_templates: dict[str, TemplateInfo] = {}
|
||
|
|
|
||
|
|
|
||
|
|
def pic_template(
|
||
|
|
name: str | None = None,
|
||
|
|
collecting: set[str] | None = None,
|
||
|
|
):
|
||
|
|
def deco(func: TemplateRenderer):
|
||
|
|
template_name = name or func.__name__
|
||
|
|
if template_name in loaded_templates:
|
||
|
|
raise ValueError(f"Template {template_name} already exists")
|
||
|
|
loaded_templates[template_name] = TemplateInfo(
|
||
|
|
renderer=func,
|
||
|
|
collectors=collecting,
|
||
|
|
)
|
||
|
|
logger.debug(f"Registered template {template_name}")
|
||
|
|
|
||
|
|
return deco
|
||
|
|
|
||
|
|
|
||
|
|
def load_builtin_templates():
|
||
|
|
for module in Path(__file__).parent.iterdir():
|
||
|
|
name = module.name
|
||
|
|
if (not module.is_dir()) or name.startswith("_"):
|
||
|
|
continue
|
||
|
|
assert importlib.import_module(f".{name}", __package__)
|
||
|
|
|
||
|
|
|
||
|
|
async def render_current_template(**kwargs: Unpack[TemplateRendererKwargs]):
|
||
|
|
return await loaded_templates[config.ps_template].renderer(**kwargs)
|