monsterxcn / nonebot-plugin-gachalogs Goto Github PK

import asyncio
from base64 import b64encode
from io import BytesIO
from typing import Dict, Literal, Tuple, Union

import matplotlib.pyplot as plt
from matplotlib import font_manager as fm
from PIL import Image

PIE_FONT = "/home/ubuntu/Service/nonebot/data/gachalogs/LXGW-Bold-minipie.ttf"  # 修改此处为自己的字体路径
stat = {
    "total": 1266,
    "cntNot5": 8,
    "cntStar3": 1068,
    "cntChar4": 123,
    "cntWeapon4": 55,
    "cntChar5": 20,
    "cntWeapon5": 0,
}


# https://github.com/monsterxcn/nonebot-plugin-gachalogs/blob/c231ab667872b10789e7f12afa6f4707dbb41b20/nonebot_plugin_gachalogs/data_render.py#L291-L344
async def drawPie(
    stat: Dict, rt: Literal["base64", "image"] = "base64"
) -> Tuple[Union[str, Image.Image], bool]:
    """
    单个饼图绘制

    * ``param stat: Dict`` 统计数据，由 ``calcStat()`` 生成
    * ``param rt: Literal["base64", "image"] = "base64"`` 返回结果类型，默认为 Base64 编码字符串，传入 ``image`` 则返回 ``PIL.Image.Image`` 对象
    - ``return: Tuple[Union[str, Image.Image], bool]`` 返回结果、是否展示三星物品数据
    """
    partMap = [
        {"label": "三星武器", "color": "#73c0de", "total": stat["cntStar3"]},
        {"label": "四星武器", "color": "#91cc75", "total": stat["cntWeapon4"]},
        {"label": "四星角色", "color": "#5470c6", "total": stat["cntChar4"]},
        {"label": "五星武器", "color": "#ee6666", "total": stat["cntWeapon5"]},
        {"label": "五星角色", "color": "#fac858", "total": stat["cntChar5"]},
    ]
    # 如果可显示项目多于 4 项，考虑隐藏三星数据
    showStar3 = True
    if len([p for p in partMap if p["total"]]) > 4:
        partMap = [p for p in partMap if p["label"] != "三星武器"]
        showStar3 = False
    # 提取数据
    labels = [p["label"] for p in partMap if p["total"]]
    colors = [p["color"] for p in partMap if p["total"]]
    sizes = [p["total"] for p in partMap if p["total"]]
    explode = [(0.05 if "五星" in p["label"] else 0) for p in partMap if p["total"]]
    # 绘制饼图
    textprops = {"fontproperties": fm.FontProperties(fname=PIE_FONT, size=18)}  # type: ignore
    fig, ax = plt.subplots()
    fig.set_facecolor("#f9f9f9")  # 这里应该可以任意调整图片底色
    ax.set_facecolor("#f9f9f9")
    ax.pie(
        sizes,
        labels=labels,
        colors=colors,
        autopct="%0.2f%%",
        labeldistance=1.1,
        pctdistance=0.7,
        startangle=60,
        radius=0.7,
        explode=explode,
        shadow=False,
        textprops=textprops,
    )
    ax.axis("equal")
    # 生成图片
    ioBytes = BytesIO()
    plt.savefig(ioBytes, format="png")
    ioBytes.seek(0)
    Image.open(ioBytes).convert("RGBA").save("pie.png")
    if rt == "base64":
        b64Code = b64encode(ioBytes.read()).decode()
        return b64Code, showStar3
    else:
        return Image.open(ioBytes), showStar3


asyncio.run(drawPie(stat, "image"))