fit nonebot.beta2

plugins/draw_card/util.py

@@ -1,34 +1,26 @@
-from nonebot.adapters.cqhttp import MessageSegment
-from typing import List, Union, Set
-from pathlib import Path
-from .config import path_dict
-from configs.path_config import IMAGE_PATH
-from utils.http_utils import AsyncHttpx
-import nonebot
-import pypinyin
-from utils.image_utils import BuildImage
 import platform
-from services.log import logger
+from asyncio.exceptions import TimeoutError
+from utils.utils import cn2py
+from utils.http_utils import AsyncHttpx
+from typing import List, Tuple, Union, Set
+from .config import draw_config, DRAW_IMAGE_PATH
+import nonebot
+from PIL import UnidentifiedImageError
+from utils.image_utils import BuildImage
+from nonebot.adapters.onebot.v11 import MessageSegment
+from nonebot.log import logger
 import random
 from dataclasses import dataclass
 import os
 import asyncio
-from PIL import UnidentifiedImageError
-try:
-    import ujson as json
-except ModuleNotFoundError:
-    import json
 
 
-driver: nonebot.Driver = nonebot.get_driver()
+driver = nonebot.get_driver()
 
 
 loop = asyncio.get_event_loop()
 
 
-headers = {'User-Agent': '"Mozilla/4.0 (compatible; MSIE 7.0; Windows NT 5.1; TencentTraveler 4.0)"'}
-
-
 @dataclass
 class BaseData:
     name: str
@@ -44,47 +36,52 @@ class UpEvent:
 
 
 async def download_img(url: str, path: str, name: str) -> bool:
-    path = path.split('_')[0]
+    path = path.split("_")[0]
     codename = cn2py(name)
-    file = Path(IMAGE_PATH + f'/draw_card/{path}/{codename}.png')
-    if not file.exists():
-        file.parent.mkdir(exist_ok=True, parents=True)
+    img_path = DRAW_IMAGE_PATH / f"{path}" / f"{codename}.png"
+    if not img_path.exists():
         try:
-            if await AsyncHttpx.download_file(url, IMAGE_PATH + f'/draw_card/{path}/{codename}.png'):
-                logger.info(f'下载 {path_dict[path]} 图片成功，名称：{name}，url：{url}')
+            if await AsyncHttpx.download_file(url, img_path):
+                logger.info(
+                    f"下载 {draw_config.path_dict[path]} 图片成功，名称：{name}，url：{url}"
+                )
                 return True
-        except Exception as e:
-            logger.warning(f'下载 {path_dict[path]} 链接错误 {type(e)}：{e}，名称：{name}，url：{url}')
+        except TimeoutError:
+            logger.warning(f"下载 {draw_config.path_dict[path]} 图片超时，名称：{name}，url：{url}")
+        # logger.info(f'{path_dict[path]} 图片 {name} 已存在')
     return False
 
 
 @driver.on_startup
 def _check_dir():
-    for dir_name in path_dict.keys():
-        _p = Path(IMAGE_PATH + f'/draw_card/' + dir_name)
-        if not _p.exists():
-            _p.mkdir(parents=True, exist_ok=True)
+    for dir_name in draw_config.path_dict.keys():
+        dir_path = DRAW_IMAGE_PATH / f"{dir_name}"
+        if not dir_path.exists():
+            dir_path.mkdir(parents=True, exist_ok=True)
 
 
-async def generate_img(card_set: Union[Set[BaseData], List[BaseData]], game_name: str, star_list: list) -> str:
+async def generate_img(
+    card_set: Union[Set[BaseData], List[BaseData]], game_name: str, star_list: list
+) -> str:
     # try:
     img_list = []
     background_list = []
     for x in card_set:
-        if game_name == 'prts':
+        if game_name == "prts":
             if x.star == 6:
-                background_list.append('#FFD700')
+                background_list.append("#FFD700")
             elif x.star == 5:
-                background_list.append('#DAA520')
+                background_list.append("#DAA520")
             elif x.star == 4:
-                background_list.append('#9370D8')
+                background_list.append("#9370D8")
             else:
-                background_list.append('white')
-        if game_name == 'azur':
-            if os.path.exists(IMAGE_PATH + f'/draw_card/{game_name}/{x.star}_star.png'):
-                background_list.append(IMAGE_PATH + f'/draw_card/{game_name}/{x.star}_star.png')
-        pyname = cn2py(x.name)
-        img_list.append(IMAGE_PATH + f'/draw_card/{game_name}/{pyname}.png')
+                background_list.append("white")
+        img_path = DRAW_IMAGE_PATH / f"{game_name}" / f"{x.star}_star.png"
+        if game_name == "azur":
+            if img_path.exists():
+                background_list.append(str(img_path))
+        py_name = cn2py(x.name)
+        img_list.append(str(DRAW_IMAGE_PATH / f"{game_name}" / f"{py_name}.png"))
     img_len = len(img_list)
     w = 100 * 10
     if img_len <= 10:
@@ -94,7 +91,9 @@ async def generate_img(card_set: Union[Set[BaseData], List[BaseData]], game_name
         h = 100 * int(img_len / 10)
     else:
         h = 100 * int(img_len / 10) + 100
-    card_img = await asyncio.get_event_loop().run_in_executor(None, _pst, h, img_list, game_name, background_list)
+    card_img = await asyncio.get_event_loop().run_in_executor(
+        None, _pst, h, img_list, game_name, background_list
+    )
     num = 0
     for n in star_list:
         num += n
@@ -108,12 +107,12 @@ def _pst(h: int, img_list: list, game_name: str, background_list: list):
     idx = 0
     for img in img_list:
         try:
-            if game_name == 'prts':
+            if game_name == "prts":
                 bk = BuildImage(100, 100, color=background_list[idx])
                 b = BuildImage(94, 94, background=img)
                 bk.paste(b, (3, 3))
                 b = bk
-            elif game_name == 'azur' and background_list:
+            elif game_name == "azur" and background_list:
                 bk = BuildImage(100, 100, background=background_list[idx])
                 b = BuildImage(98, 90, background=img)
                 bk.paste(b, (1, 5))
@@ -122,38 +121,75 @@ def _pst(h: int, img_list: list, game_name: str, background_list: list):
                 try:
                     b = BuildImage(100, 100, background=img)
                 except UnidentifiedImageError as e:
-                    logger.warning(f'无法识别图片 已删除图片，下次更新重新下载... e：{e}')
+                    logger.warning(f"无法识别图片 已删除图片，下次更新重新下载... e：{e}")
                     if os.path.exists(img):
                         os.remove(img)
-                    b = BuildImage(100, 100, color='black')
+                    b = BuildImage(100, 100, color="black")
         except FileNotFoundError:
-            logger.warning(f'{img} not exists')
-            b = BuildImage(100, 100, color='black')
+            logger.warning(f"{img} not exists")
+            b = BuildImage(100, 100, color="black")
         card_img.paste(b)
         idx += 1
     return card_img
 
 
-def init_star_rst(star_list: list, cnlist: list, max_star_list: list, max_star_index_list: list, up_list: list = None) -> str:
+# 初始化输出数据
+def init_star_rst(
+    star_list: list,
+    cnlist: list,
+    max_star_list: list,
+    max_star_index_list: list,
+    up_list: list = None,
+) -> str:
     if not up_list:
         up_list = []
-    rst = ''
+    rst = ""
     for i in range(len(star_list)):
         if star_list[i]:
-            rst += f'[{cnlist[i]}×{star_list[i]}] '
-    rst += '\n'
+            rst += f"[{cnlist[i]}×{star_list[i]}] "
+    rst += "\n"
     for i in range(len(max_star_list)):
         if max_star_list[i] in up_list:
-            rst += f'第 {max_star_index_list[i]+1} 抽获取UP {max_star_list[i]}\n'
+            rst += f"第 {max_star_index_list[i]+1} 抽获取UP {max_star_list[i]}\n"
         else:
-            rst += f'第 {max_star_index_list[i]+1} 抽获取 {max_star_list[i]}\n'
+            rst += f"第 {max_star_index_list[i]+1} 抽获取 {max_star_list[i]}\n"
     return rst
 
 
+# 更好的初始化
+def init_rst(
+    max_star_char_dict: dict,
+    star_num_list: List[int],
+    star: List[str],
+    up_list: list = None,
+):
+    # print(max_star_char_dict)
+    # print(star_num_list)
+    # print(up_list)
+    up_list = up_list if up_list else []
+    rst = ""
+    for i in range(len(star_num_list)):
+        if star_num_list[i]:
+            rst += f"[{star[i]}×{star_num_list[i]}] "
+    rst += "\n"
+    _tmp = []
+    for name in max_star_char_dict.keys():
+        _tmp += max_star_char_dict[name]
+    for index in sorted(_tmp):
+        for name in max_star_char_dict.keys():
+            if index in max_star_char_dict[name]:
+                if name in up_list:
+                    rst += f"第 {index} 抽获取UP {name}\n"
+                else:
+                    rst += f"第 {index} 抽获取 {name}\n"
+    print(rst)
+    return rst[:-1] if rst else ""
+
+
 def max_card(_dict: dict):
     _max_value = max(_dict.values())
     _max_user = list(_dict.keys())[list(_dict.values()).index(_max_value)]
-    return f'抽取到最多的是{_max_user}，共抽取了{_max_value}次'
+    return f"抽取到最多的是{_max_user}，共抽取了{_max_value}次"
     # ThreeHighest = nlargest(3, operator_dict, key=operator_dict.get)
     # rst = '最喜欢你的前三位是干员是：\n'
     # for name in ThreeHighest:
@@ -161,6 +197,55 @@ def max_card(_dict: dict):
     # return rst[:-1]
 
 
+# 获取up和概率
+async def init_up_char(announcement):
+    UP_CHAR = []
+    UP_ARMS = []
+    tmp = ""
+    up_char_dict = await announcement.update_up_char()
+    for x in list(up_char_dict.keys()):
+        tmp += up_char_dict[x]["title"] + "[\n]"
+    tmp = tmp.split("[\n]")
+    _CURRENT_CHAR_POOL_TITLE = tmp[0]
+    if len(up_char_dict) > 1:
+        _CURRENT_ARMS_POOL_TITLE = tmp[1]
+    else:
+        _CURRENT_ARMS_POOL_TITLE = ""
+    POOL_IMG = ""
+    x = [x for x in list(up_char_dict.keys())]
+    if _CURRENT_CHAR_POOL_TITLE:
+        POOL_IMG += MessageSegment.image(up_char_dict[x[0]]["pool_img"])
+    try:
+        if _CURRENT_ARMS_POOL_TITLE:
+            POOL_IMG += MessageSegment.image(up_char_dict[x[1]]["pool_img"])
+    except (IndexError, KeyError):
+        pass
+    logger.info(
+        f"成功获取{announcement.game_name}当前up信息...当前up池: {_CURRENT_CHAR_POOL_TITLE} & {_CURRENT_ARMS_POOL_TITLE}"
+    )
+    for key in up_char_dict.keys():
+        for star in up_char_dict[key]["up_char"].keys():
+            up_char_lst = []
+            for char in up_char_dict[key]["up_char"][star].keys():
+                up_char_lst.append(char)
+            if up_char_lst:
+                if key == "char":
+                    UP_CHAR.append(
+                        UpEvent(star=int(star), operators=up_char_lst, zoom=0)
+                    )
+                else:
+                    UP_ARMS.append(
+                        UpEvent(star=int(star), operators=up_char_lst, zoom=0)
+                    )
+    return (
+        _CURRENT_CHAR_POOL_TITLE,
+        _CURRENT_ARMS_POOL_TITLE,
+        POOL_IMG,
+        UP_CHAR,
+        UP_ARMS,
+    )
+
+
 def is_number(s) -> bool:
     try:
         float(s)
@@ -169,6 +254,7 @@ def is_number(s) -> bool:
         pass
     try:
         import unicodedata
+
         unicodedata.numeric(s)
         return True
     except (TypeError, ValueError):
@@ -176,13 +262,6 @@ def is_number(s) -> bool:
     return False
 
 
-def cn2py(word) -> str:
-    temp = ""
-    for i in pypinyin.pinyin(word, style=pypinyin.NORMAL):
-        temp += ''.join(i)
-    return temp
-
-
 def set_list(lst: List[BaseData]) -> list:
     tmp = []
     name_lst = []
@@ -193,6 +272,7 @@ def set_list(lst: List[BaseData]) -> list:
     return tmp
 
 
+# 获取星级
 def get_star(star_lst: List[int], probability_lst: List[float]) -> int:
     rand = random.random()
     add = 0
@@ -206,49 +286,14 @@ def get_star(star_lst: List[int], probability_lst: List[float]) -> int:
             return star_lst[i]
 
 
-# 获取up和概率
-async def init_up_char(announcement):
-    UP_CHAR = []
-    UP_ARMS = []
-    tmp = ''
-    up_char_dict = await announcement.update_up_char()
-    for x in list(up_char_dict.keys()):
-        tmp += up_char_dict[x]['title'] + '[\n]'
-    tmp = tmp.split('[\n]')
-    _CURRENT_CHAR_POOL_TITLE = tmp[0]
-    if len(up_char_dict) > 1:
-        _CURRENT_ARMS_POOL_TITLE = tmp[1]
-    else:
-        _CURRENT_ARMS_POOL_TITLE = ''
-    POOL_IMG = ''
-    x = [x for x in list(up_char_dict.keys())]
-    if _CURRENT_CHAR_POOL_TITLE:
-        POOL_IMG += MessageSegment.image(up_char_dict[x[0]]['pool_img'])
-    try:
-        if _CURRENT_ARMS_POOL_TITLE:
-            POOL_IMG += MessageSegment.image(up_char_dict[x[1]]['pool_img'])
-    except (IndexError, KeyError):
-        pass
-    logger.info(f'成功获取{announcement.game_name}当前up信息...当前up池: {_CURRENT_CHAR_POOL_TITLE} & {_CURRENT_ARMS_POOL_TITLE}')
-    for key in up_char_dict.keys():
-        for star in up_char_dict[key]['up_char'].keys():
-            up_char_lst = []
-            for char in up_char_dict[key]['up_char'][star].keys():
-                up_char_lst.append(char)
-            if up_char_lst:
-                if key == 'char':
-                    UP_CHAR.append(UpEvent(star=int(star), operators=up_char_lst, zoom=0))
-                else:
-                    UP_ARMS.append(UpEvent(star=int(star), operators=up_char_lst, zoom=0))
-    return _CURRENT_CHAR_POOL_TITLE, _CURRENT_ARMS_POOL_TITLE, POOL_IMG, UP_CHAR, UP_ARMS
-
-
 # 整理数据
-def format_card_information(count: int, star_list: List[int], func, pool_name: str = '', guaranteed: bool = True):
-    max_star_lst = []       # 获取的最高星级角色列表
-    max_index_lst = []      # 获取最高星级角色的次数
-    obj_list = []           # 获取所有角色
-    obj_dict = {}           # 获取角色次数字典
+def format_card_information(
+    count: int, star_list: List[int], func, pool_name: str = "", guaranteed: bool = True
+):
+    max_star_lst = []  # 获取的最高星级角色列表
+    max_index_lst = []  # 获取最高星级角色的次数
+    obj_list = []  # 获取所有角色
+    obj_dict = {}  # 获取角色次数字典
     _count = -1
     if guaranteed:
         _count = 0
@@ -285,28 +330,28 @@ def format_card_information(count: int, star_list: List[int], func, pool_name: s
 
 
 # 检测次数是否合法
-def check_num(num: str, max_num: int) -> 'str, bool':
+def check_num(num: str, max_num: int) -> Tuple[str, bool]:
     if is_number(num):
         try:
             num = int(num)
         except ValueError:
-            return '必！须！是！数！字！', False
+            return "必！须！是！数！字！", False
     if num > max_num:
-        return '一井都满不足不了你嘛！快爬开！', False
+        return "一井都满不足不了你嘛！快爬开！", False
     if num < 1:
-        return '虚空抽卡？？？', False
+        return "虚空抽卡？？？", False
     else:
         return str(num), True
 
 
 # 移除windows和linux下特殊字符
-def remove_prohibited_str(name: str):
-    if platform.system().lower() == 'windows':
-        tmp = ''
+def remove_prohibited_str(name: str) -> str:
+    if platform.system().lower() == "windows":
+        tmp = ""
         for i in name:
-            if i not in ['\\', '/', ':', '*', '?', '"', '<', '>', '|']:
+            if i not in ["\\", "/", ":", "*", "?", '"', "<", ">", "|"]:
                 tmp += i
         name = tmp
     else:
-        name = name.replace('/', '\\')
+        name = name.replace("/", "\\")
     return name