From 6a1333ebb51812627b879816ab9750a0db8c87d1 Mon Sep 17 00:00:00 2001
From: xtaodada <xtao@xtaolink.cn>
Date: Fri, 28 Apr 2023 19:32:53 +0800
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20support:=20square=20icon?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 func/fetch_avatars.py       |  1 +
 models/avatar.py            |  6 ++++--
 models/avatar_config.py     | 18 +++++++++++++++++-
 models/light_cone.py        |  4 ++--
 models/light_cone_config.py |  8 ++++++++
 models/material.py          |  2 +-
 models/monster.py           |  2 +-
 res_func/avatar.py          | 13 ++++++++-----
 8 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/func/fetch_avatars.py b/func/fetch_avatars.py
index f7b6246..5934f96 100644
--- a/func/fetch_avatars.py
+++ b/func/fetch_avatars.py
@@ -27,6 +27,7 @@ async def fetch_avatars(data: Children):
         avatar = Avatar(
             id=content.content_id,
             name=content.title,
+            icon=content.icon,
             quality=m_quality,
             element=m_element,
             destiny=m_destiny,
diff --git a/models/avatar.py b/models/avatar.py
index f069d08..51759fa 100644
--- a/models/avatar.py
+++ b/models/avatar.py
@@ -1,7 +1,7 @@
 from typing import List
 from pydantic import BaseModel
-from models.enums import Quality, Destiny, Element
-from models.material import Material
+from .enums import Quality, Destiny, Element
+from .material import Material
 
 
 class AvatarInfo(BaseModel):
@@ -44,6 +44,8 @@ class Avatar(BaseModel):
     """角色ID"""
     name: str
     """名称"""
+    icon: str
+    """图标"""
     quality: Quality
     """品质"""
     destiny: Destiny
diff --git a/models/avatar_config.py b/models/avatar_config.py
index 19d77d2..87c231b 100644
--- a/models/avatar_config.py
+++ b/models/avatar_config.py
@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 from pydantic import BaseModel
 
@@ -22,3 +22,19 @@ class AvatarIcon(BaseModel):
     """名称"""
     icon: List[str]
     """图标（从小到大）"""
+
+    @property
+    def gacha(self) -> str:
+        return self.icon[3]
+
+    @property
+    def icon_(self) -> str:
+        return self.icon[0]
+
+    @property
+    def square(self) -> Optional[str]:
+        return self.icon[1]
+
+    @property
+    def normal(self) -> str:
+        return self.icon[2]
diff --git a/models/light_cone.py b/models/light_cone.py
index d134a5e..0cb5a60 100644
--- a/models/light_cone.py
+++ b/models/light_cone.py
@@ -1,8 +1,8 @@
 # 光锥
 from pydantic import BaseModel
 
-from models.enums import Quality, Destiny
-from models.material import Material
+from .enums import Quality, Destiny
+from .material import Material
 
 
 class LightConeItem(BaseModel):
diff --git a/models/light_cone_config.py b/models/light_cone_config.py
index f8698c2..88fb96f 100644
--- a/models/light_cone_config.py
+++ b/models/light_cone_config.py
@@ -10,3 +10,11 @@ class LightConeIcon(BaseModel):
     """名称"""
     icon: List[str]
     """图标（从小到大）"""
+
+    @property
+    def gacha(self) -> str:
+        return self.icon[1]
+
+    @property
+    def icon_(self) -> str:
+        return self.icon[0]
diff --git a/models/material.py b/models/material.py
index d440ca4..88ab22f 100644
--- a/models/material.py
+++ b/models/material.py
@@ -1,7 +1,7 @@
 # 材料
 from pydantic import BaseModel
 
-from models.enums import Quality, MaterialType
+from .enums import Quality, MaterialType
 
 
 class Material(BaseModel):
diff --git a/models/monster.py b/models/monster.py
index 68f4551..ccc20d6 100644
--- a/models/monster.py
+++ b/models/monster.py
@@ -1,7 +1,7 @@
 # 敌对物种
 from pydantic import BaseModel
 
-from models.enums import MonsterType, Area
+from .enums import MonsterType, Area
 
 
 class Monster(BaseModel):
diff --git a/res_func/avatar.py b/res_func/avatar.py
index b65597e..2b63ac8 100644
--- a/res_func/avatar.py
+++ b/res_func/avatar.py
@@ -28,18 +28,21 @@ async def fetch_config(text_map: Dict[str, str]) -> List[AvatarConfig]:
     return datas
 
 
-async def parse_station(datas, tag: Tag, avatar: AvatarConfig):
-    second_pic = f'{base_station_url}{tag.find("img").get("src")}'
+async def parse_station(datas, name: str, tag: Tag, avatar: AvatarConfig):
+    second_pic = ""
+    if avatar_model := all_avatars_name.get(name):
+        second_pic = avatar_model.icon
+    third_pic = f'{base_station_url}{tag.find("img").get("src")}'
     html = await client.get(f'{base_station_url}{tag.get("href")}')
     soup = BeautifulSoup(html.text, "lxml")
     text = soup.find("div", {"class": "a6678 a4af5"}).get("style")
-    third_pic = f'{base_station_url}{text[text.find("(") + 2:text.find(")") - 1]}' if text else ""
+    four_pic = f'{base_station_url}{text[text.find("(") + 2:text.find(")") - 1]}' if text else ""
     first_pic = f'{base_station_url}{soup.find("img", {"class": "ac39b a6602"}).get("src")}'
     datas.append(
         AvatarIcon(
             id=avatar.AvatarID,
             name=avatar.name,
-            icon=[first_pic, second_pic, third_pic],
+            icon=[first_pic, second_pic, third_pic, four_pic],
         )
     )
 
@@ -64,7 +67,7 @@ async def fetch_station(configs_map: Dict[str, AvatarConfig]) -> List[AvatarIcon
         if name == "开拓者" and ktz:
             continue
         if avatar_model := configs_map.get(name):
-            tasks.append(parse_station(datas, avatar, avatar_model))
+            tasks.append(parse_station(datas, name, avatar, avatar_model))
             if name == "开拓者":
                 ktz = True
         else: