✨🎨 完善插件, 优化代码

1 files changed, 181 insertions, 35 deletions

diff --git a/ATRI/plugins/anti_effort/data_source.py b/ATRI/plugins/anti_effort/data_source.py
index 0a0ef02..ba9cdf5 100644
--- a/ATRI/plugins/anti_effort/data_source.py
+++ b/ATRI/plugins/anti_effort/data_source.py
@@ -1,16 +1,23 @@
-import json
 import re
 import os
+import json
 from pathlib import Path
+from datetime import datetime
+from tabulate import tabulate
 
-from ATRI.service import Service
+from ATRI import driver
+from ATRI.service import Service, ServiceTools
 from ATRI.rule import is_in_service
 from ATRI.utils import request
 from ATRI.log import logger as log
 
-from .models import AntiEffortUserModel
+from .image_dealer import image_dealer
+from .models import AntiEffortModel, AntiEffortUserModel
 
 
+TENCENT_AVATER_URL = "https://q1.qlogo.cn/g?b=qq&nk={user_id}&s=640"
+SOURCE_URL = "https://fastly.jsdelivr.net/gh/Kyomotoi/CDN@master/project/ATRI/"
+
 PLUGIN_DIR = Path(".") / "data" / "plugins" / "anti_effort"
 PLUGIN_DIR.mkdir(parents=True, exist_ok=True)
 
@@ -21,10 +28,42 @@ class AntiEffort(Service):
             self, "谁是卷王", "谁是卷王！", rule=is_in_service("谁是卷王"), main_cmd="/ae"
         )
 
-    def add_user(
+    def get_enabled_group(self) -> list:
+        groups = list()
+
+        files = os.listdir(PLUGIN_DIR)
+        if not files:
+            return groups
+
+        for f in files:
+            raw_data = f.split(".")
+            if raw_data[-1] != "json":
+                continue
+
+            if "-ld" in raw_data[0]:
+                patt = r"([0-9].*?)-ld"
+                match = re.findall(patt, raw_data[0])
+                if match:
+                    group_id = int(match[0])
+                    groups.append(group_id)
+            else:
+                group_id = int(raw_data[0])
+                groups.append(group_id)
+
+        return groups
+
+    def get_data(self, group_id: int) -> dict:
+        file_path = PLUGIN_DIR / f"{group_id}.json"
+        if not file_path.exists():
+            return dict()
+        return json.loads(file_path.read_text())
+
+    async def add_user(
         self, group_id: int, user_id: int, user_nickname: str, waka_url: str
     ) -> str:
-        patt = r"/@([a-zA-Z0-9].*?)/([a-zA-Z0-9].*?).json"
+        patt = (
+            r"https:\/\/wakatime.com\/share\/@([a-zA-Z0-9].*?)\/([a-zA-Z0-9].*?).json"
+        )
         match = re.findall(patt, waka_url)
         if not match:
             return "哥，你提供的链接有问题啊"
@@ -33,74 +72,100 @@ class AntiEffort(Service):
 
         file_path = PLUGIN_DIR / f"{group_id}.json"
         if not file_path.exists():
+            now_time = datetime.now().strftime("%Y/%m/%d - %H:%M:%S")
+            d = AntiEffortModel(update_time=now_time, data=list()).dict()
             with open(file_path, "w") as f:
-                f.write(json.dumps(list()))
+                f.write(json.dumps(d))
 
-        data: list = json.loads(file_path.read_text())
+        raw_data = json.loads(file_path.read_text())
+        data: list = raw_data["data"]
         for i in data:
             if i["user_id"] == user_id:
                 return "你已经在卷王统计榜力！"
 
+        try:
+            resp = await request.get(waka_url)
+            user_w_data = resp.json()  # type: ignore
+
+            d = user_w_data["data"]
+            last_7_days_count = float()
+            recent_count = float(d[-1]["grand_total"]["decimal"])
+            for u in d:
+                last_7_days_count += float(u["grand_total"]["decimal"])
+        except Exception:
+            log.error(f"获取卷王 {w_user_id} 数据失败！")
+            last_7_days_count = int()
+            recent_count = int()
+
         data.append(
             AntiEffortUserModel(
                 user_id=user_id,
                 user_nickname=user_nickname,
                 w_user_id=w_user_id,
                 waka_url=waka_url,
-                last_7_days_count=int(),
-                recent_count=int(),
+                last_7_days_count=last_7_days_count,
+                recent_count=recent_count,
             ).dict()
         )
+        raw_data["data"] = data
         with open(file_path, "w") as f:
-            f.write(json.dumps(data))
+            f.write(json.dumps(raw_data))
 
         return "成功加入卷王统计榜！"
 
     def del_user(self, group_id: int, user_id: int) -> str:
-        data = self.get_data(group_id)
-        file_path = PLUGIN_DIR / f"{group_id}.json"
-        if file_path.exists():
+        raw_data = self.get_data(group_id)
+        if raw_data:
+            file_path = PLUGIN_DIR / f"{group_id}.json"
+            data = raw_data["data"]
             for i in data:
                 if i["user_id"] == user_id:
                     data.remove(i)
-                    with open(PLUGIN_DIR / f"{group_id}.json", "w") as f:
-                        f.write(json.dumps(data))
+                    raw_data["data"] = data
+                    with open(file_path, "w") as f:
+                        f.write(json.dumps(raw_data))
                     return "成功退出卷王统计榜！"
 
         return "你未加入卷王统计榜捏"
 
-    def get_data(self, group_id: int) -> list:
-        file_path = PLUGIN_DIR / f"{group_id}.json"
-        if not file_path.exists():
-            return list()
-        return json.loads(file_path.read_text())
-
     def update_user(self, group_id: int, user_id: int, update_map: dict):
-        data = self.get_data(group_id)
-        if data:
+        raw_data = self.get_data(group_id)
+        if raw_data:
+            now_time = datetime.now().strftime("%Y/%m/%d - %H:%M:%S")
             file_path = PLUGIN_DIR / f"{group_id}.json"
+            data = raw_data["data"]
             for i in data:
                 if i["user_id"] == user_id:
                     for k, v in update_map.items():
                         i[k] = v
+                    raw_data["update_time"] = now_time
+                    raw_data["data"] = data
                     with open(file_path, "w") as f:
-                        f.write(json.dumps(data))
+                        f.write(json.dumps(raw_data))
 
-    async def update_data(self):
-        files = os.listdir(PLUGIN_DIR)
-        if not files:
-            return 114514
+    def store_user_data_recent(self):
+        groups = self.get_enabled_group()
+        if not groups:
+            return
 
-        groups = list()
-        for f in files:
-            group_id = int(f.split(".")[0])
-            groups.append(group_id)
+        for g in groups:
+            data = self.get_data(g)
+            if not data:
+                return
+
+            file_path = PLUGIN_DIR / f"{g}-ld.json"
+            with open(file_path, "w") as f:
+                f.write(json.dumps(data))
+
+    async def update_data(self):
+        groups = self.get_enabled_group()
 
         for i in groups:
-            data = self.get_data(i)
-            if not data:
+            raw_data = self.get_data(i)
+            if not raw_data:
                 continue
 
+            data = raw_data["data"]
             for j in data:
                 try:
                     resp = await request.get(j["waka_url"])
@@ -109,7 +174,10 @@ class AntiEffort(Service):
                     log.error(f"获取卷王 {j['user_id']} 数据失败！")
                     continue
 
-                d = user_w_data["data"]
+                d = user_w_data.get("data", dict())
+                if not d:
+                    continue
+
                 user_w_last_7_days_count = float()
                 user_w_recent_count = float(d[-1]["grand_total"]["decimal"])
                 for u in d:
@@ -120,3 +188,81 @@ class AntiEffort(Service):
                     "recent_count": user_w_recent_count,
                 }
                 self.update_user(i, j["user_id"], u_data)
+
+    async def gen_img(
+        self, user_id: int, user_nickname: str, coding_time: float
+    ) -> str:
+
+        try:
+            resp = await request.get(TENCENT_AVATER_URL.format(user_id=user_id))
+        except Exception:
+            log.error("插件 anti_effort 获取用户头像失败！")
+            return str()
+
+        avatar_byt = resp.read()  # type: ignore
+        result = image_dealer(
+            avatar_byt, user_nickname, str(round(coding_time, 2)) + " hrs"
+        )
+        return f"file:///{result}"
+
+    def gen_rank(self, raw_data: dict, user_id: int, typ: str) -> str:
+        table_type = "Today"
+        sort_type = "recent_count"
+        if typ == "today":
+            rank_type = "今日"
+        elif typ == "recent_week":
+            rank_type = "近一周"
+            table_type = "Last 7 Days"
+            sort_type = "last_7_days_count"
+        elif typ == "global_today":
+            rank_type = "今日公共"
+        else:
+            rank_type = "近一周公共"
+            table_type = "Last 7 Days"
+            sort_type = "last_7_days_count"
+
+        data = raw_data["data"]
+        data = sorted(data, key=lambda x: x[sort_type], reverse=True)
+
+        user_rank = 0
+        for i, user in enumerate(data):
+            if user["user_id"] == user_id:
+                user_rank = i + 1
+                break
+
+        table = [
+            [
+                f"{i + 1}",
+                f"{x['user_nickname']}",
+                f"{round(x['recent_count'], 2)} hrs",
+            ]
+            for i, x in enumerate(data)
+        ][:10]
+        table.insert(0, ["Rank", "Member", table_type])
+        result = tabulate(table, tablefmt="plain")
+        update_time = raw_data["update_time"]
+        rank = f"\n你位于第 {user_rank} 名" if user_rank else str()
+        repo = f"《{rank_type}十佳卷王》\nUpdate Time: {update_time}\n{result}{rank}"
+        return repo
+
+
+async def init_source():
+    files = ["xb-bg-0.png", "hwxw.ttf"]
+
+    try:
+        for i in files:
+            file_path = PLUGIN_DIR / i
+            if not file_path.is_file():
+                log.warning("插件 anti_effort 缺少所需资源，装载中")
+                url = SOURCE_URL + i
+                data = await request.get(url)
+                with open(file_path, "wb") as f:
+                    f.write(data.read())  # type: ignore
+    except Exception:
+        ServiceTools.service_controller("谁是卷王", False)
+        log.error(f"插件 anti_effort 装载资源失败. 已自动禁用")
+
+    log.success("插件 anti_effort 装载资源完成")
+
+
+driver().on_startup(init_source)