From 9e87a1983b7742e75f801b08f1ce2873fce029bb Mon Sep 17 00:00:00 2001
From: xtaodada <xtao@xtaolink.cn>
Date: Sun, 26 May 2024 19:46:32 +0800
Subject: [PATCH] :bug: Fix weapon icon

---
 func/base.py      |  7 ++++++-
 func/character.py |  2 +-
 func/weapon.py    | 23 ++++++++++++++++++++++-
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/func/base.py b/func/base.py
index 016eea2..7250fef 100644
--- a/func/base.py
+++ b/func/base.py
@@ -15,7 +15,7 @@
 async def get_item_base(item_path: str, item_type: Type[BaseModel], fix_func=None):
     req = await client.get(f"{BASE_URL}/{item_path}")
     req.raise_for_status()
-    data = fix_func(req.json()) if fix_func else req.json()
+    data = await fix_func(req.json()) if fix_func else req.json()
     item_data = item_type(**data)
     return item_data
 
@@ -25,3 +25,8 @@ def save_data(item_path: str, item_datas: List[BaseModel]):
     item_data = [i.dict(by_alias=True) for i in item_datas]
     with path.open("w", encoding="utf-8") as f:
         json.dump(item_data, f, ensure_ascii=False, indent=4)
+
+
+async def have_url(url: str) -> bool:
+    req = await client.head(url)
+    return req.status_code == 200
diff --git a/func/character.py b/func/character.py
index 8d2bbd9..165d03a 100644
--- a/func/character.py
+++ b/func/character.py
@@ -8,7 +8,7 @@ async def get_characters():
     return await get_item_base("character", EncoreAvatars)
 
 
-def fix_character(data: Dict[str, Any]) -> Dict[str, Any]:
+async def fix_character(data: Dict[str, Any]) -> Dict[str, Any]:
     data["Name"] = data["Name"]["Content"]
     return data
 
diff --git a/func/weapon.py b/func/weapon.py
index dfadac8..d6da9dd 100644
--- a/func/weapon.py
+++ b/func/weapon.py
@@ -1,3 +1,5 @@
+from typing import Dict, Any
+
 from .base import get_item_base, save_data
 from models.weapon import EncoreWeapon, EncoreWeapons
 
@@ -6,8 +8,27 @@ async def get_weapons():
     return await get_item_base("weapon", EncoreWeapons)
 
 
+def get_url_ext(url: str) -> str:
+    return url.split(".")[-1]
+
+
+def replace_url_ext(url: str, ext: str) -> str:
+    return ".".join(url.split(".")[:-1]) + f".{ext}"
+
+
+async def fix_weapon_icon(data: Dict[str, Any]) -> Dict[str, Any]:
+    for k in ["Icon", "IconMiddle", "IconSmall"]:
+        if data[k] and data[k].startswith("/"):
+            data[k] = f"https://api.encore.moe/resource/Data{data[k]}"
+        if get_url_ext(data[k]) == "png":
+            continue
+        new_url = replace_url_ext(data[k], "png")
+        data[k] = new_url
+    return data
+
+
 async def get_weapon(weapon_id: int):
-    return await get_item_base(f"weapon/{weapon_id}", EncoreWeapon)
+    return await get_item_base(f"weapon/{weapon_id}", EncoreWeapon, fix_func=fix_weapon_icon)
 
 
 async def main():