imp: Use json_data to save ChallengeHistory

Author: seriaati

hoyo_buddy/db/models.py

@@ -6,7 +6,7 @@
 import datetime
 import pickle
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import genshin
 import orjson
@@ -30,6 +30,8 @@
 from ..utils import blur_uid, get_now
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping
+
     import aiohttp
 
     from ..hoyo.clients.gpy import GenshinClient
@@ -335,10 +337,11 @@ class ChallengeHistory(BaseModel):
     season_id = fields.IntField()
     name: fields.Field[str | None] = fields.CharField(max_length=64, null=True)
     challenge_type = fields.CharEnumField(ChallengeType, max_length=32)
-    data = fields.BinaryField()
+    data: fields.Field[bytes | None] = fields.BinaryField(null=True)
     start_time = fields.DatetimeField()
     end_time = fields.DatetimeField()
     lang: fields.Field[str | None] = fields.CharField(max_length=5, null=True)
+    json_data: fields.Field[dict[str, Any] | None] = fields.JSONField(null=True)
 
     class Meta:
         unique_together = ("uid", "season_id", "challenge_type")
@@ -353,17 +356,50 @@ def duration_str(self) -> str:
     @property
     def parsed_data(self) -> ChallengeWithLang:
         """Parsed challenge data from binary pickled data."""
-        challenge = pickle.loads(self.data)
+        if self.json_data is None:
+            if self.data is None:
+                # This shouldn't happen, data could be None because of migration to use json_data
+                msg = "Both json_data and data are None in ChallengeHistory"
+                raise ValueError(msg)
+            challenge = pickle.loads(self.data)
+        else:
+            challenge = self.load_data(self.json_data, challenge_type=self.challenge_type)
+
         lang = getattr(challenge, "lang", None)
         if lang is None:
             # NOTE: Backward compatibility, old data has lang attr, new data doesn't
             challenge.__dict__["lang"] = self.lang
-        return challenge
+        return cast("ChallengeWithLang", challenge)
+
+    @classmethod
+    def load_data(cls, raw: Mapping[str, Any], *, challenge_type: ChallengeType) -> Challenge:
+        if challenge_type is ChallengeType.SPIRAL_ABYSS:
+            return genshin.models.SpiralAbyss(**raw)
+        if challenge_type is ChallengeType.IMG_THEATER:
+            return genshin.models.ImgTheaterData(**raw)
+        if challenge_type is ChallengeType.SHIYU_DEFENSE:
+            return genshin.models.ShiyuDefense(**raw)
+        if challenge_type is ChallengeType.ASSAULT:
+            return genshin.models.DeadlyAssault(**raw)
+        if challenge_type is ChallengeType.APC_SHADOW:
+            return genshin.models.StarRailAPCShadow(**raw)
+        if challenge_type is ChallengeType.MOC:
+            return genshin.models.StarRailChallenge(**raw)
+        if challenge_type is ChallengeType.PURE_FICTION:
+            return genshin.models.StarRailPureFiction(**raw)
 
     @classmethod
     async def add_data(
-        cls, *, uid: int, challenge_type: ChallengeType, season_id: int, data: Challenge, lang: str
+        cls,
+        *,
+        uid: int,
+        challenge_type: ChallengeType,
+        season_id: int,
+        raw: Mapping[str, Any],
+        lang: str,
     ) -> None:
+        data = cls.load_data(raw, challenge_type=challenge_type)
+
         if isinstance(data, genshin.models.SpiralAbyss | genshin.models.DeadlyAssault):
             start_time = data.start_time
             end_time = data.end_time
@@ -390,15 +426,15 @@ async def add_data(
                 uid=uid,
                 season_id=season_id,
                 challenge_type=challenge_type,
-                data=pickle.dumps(data),
                 start_time=start_time,
                 end_time=end_time,
                 name=name,
                 lang=lang,
+                json_data=raw,
             )
         except exceptions.IntegrityError:
             await cls.filter(uid=uid, season_id=season_id, challenge_type=challenge_type).update(
-                data=pickle.dumps(data), name=name, lang=lang
+                name=name, lang=lang, json_data=raw
             )
 
 

hoyo_buddy/draw/funcs/hoyo/zzz/shiyu.py

@@ -224,8 +224,14 @@ def draw_frontiers(self, im: Image.Image, drawer: Drawer) -> None:
                             style="bold",
                             anchor="mm",
                         )
+
+                        # Backward compatibility, ShiyuDefenseCharacter.mindscape is added in
+                        # https://github.com/thesadru/genshin.py/commit/4e17d37f84048d2b0a478b45e374f980a7bbe3a3
+                        rank: int = getattr(agent, "mindscape", None) or self.agent_ranks.get(
+                            agent.id, 0
+                        )
                         drawer.write(
-                            str(self.agent_ranks.get(agent.id, 0)),
+                            str(rank),
                             size=36,
                             position=(start_pos[0] + 158, start_pos[1] + 22),
                             color=self.white,

hoyo_buddy/ui/hoyo/challenge.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 import discord
 from ambr.utils import remove_html_tags
@@ -19,6 +19,7 @@
     TheaterBuff,
 )
 from genshin.models import Character as GICharacter
+from loguru import logger
 
 from hoyo_buddy.bot.error_handler import get_error_embed
 from hoyo_buddy.constants import GAME_CHALLENGE_TYPES, GPY_LANG_TO_LOCALE, TRAVELER_IDS
@@ -218,47 +219,74 @@ async def _fetch_data(self) -> None:
         client = self.account.client
         client.set_lang(self.locale)
 
-        if (
-            self.challenge_type in {ChallengeType.SPIRAL_ABYSS, ChallengeType.IMG_THEATER}
-            and not self.characters
-        ):
-            self.characters = await client.get_genshin_characters(self.account.uid)
-
         await client.get_record_cards()
 
         for previous in (False, True):
             if self.challenge_type is ChallengeType.SPIRAL_ABYSS:
-                challenge = await client.get_genshin_spiral_abyss(
-                    self.account.uid, previous=previous
+                raw = await client.get_genshin_spiral_abyss(
+                    self.account.uid, previous=previous, raw=True
                 )
             elif self.challenge_type is ChallengeType.MOC:
-                challenge = await client.get_starrail_challenge(self.account.uid, previous=previous)
+                raw = await client.get_starrail_challenge(
+                    self.account.uid, previous=previous, raw=True
+                )
             elif self.challenge_type is ChallengeType.PURE_FICTION:
-                challenge = await client.get_starrail_pure_fiction(
-                    self.account.uid, previous=previous
+                raw = await client.get_starrail_pure_fiction(
+                    self.account.uid, previous=previous, raw=True
                 )
             elif self.challenge_type is ChallengeType.APC_SHADOW:
-                challenge = await client.get_starrail_apc_shadow(
-                    self.account.uid, previous=previous
+                raw = await client.get_starrail_apc_shadow(
+                    self.account.uid, previous=previous, raw=True
                 )
             elif self.challenge_type is ChallengeType.IMG_THEATER:
-                challenges = (
-                    await client.get_imaginarium_theater(self.account.uid, previous=previous)
-                ).datas
-                if not challenges:
+                raw_ = await client.get_imaginarium_theater(
+                    self.account.uid, previous=previous, raw=True
+                )
+                datas: list[dict[str, Any]] = raw_.get("data", [])
+                if not datas:
                     raise NoChallengeDataError(ChallengeType.IMG_THEATER)
 
-                challenge = max(challenges, key=lambda c: c.stats.difficulty.value)
+                try:
+                    raw = max(datas, key=lambda d: d["stat"]["difficulty_id"])
+                except KeyError:
+                    logger.error("Failed to get max difficulty ID from data", datas=datas)
+                    raw = datas[-1]
             elif self.challenge_type is ChallengeType.SHIYU_DEFENSE:
-                if not self.agent_ranks:
+                raw = await client.get_shiyu_defense(self.account.uid, previous=previous, raw=True)
+                challenge = ChallengeHistory.load_data(raw, challenge_type=self.challenge_type)
+                challenge = cast("ShiyuDefense", challenge)
+
+                # Backward compatibility, ShiyuDefenseCharacter.mindscape is added in
+                # https://github.com/thesadru/genshin.py/commit/4e17d37f84048d2b0a478b45e374f980a7bbe3a3
+                is_new_ver = (
+                    challenge.floors
+                    and challenge.floors[0].node_1.characters
+                    and hasattr(challenge.floors[0].node_1.characters[0], "mindscape")
+                )
+
+                # No need to fetch agent ranks if the data is using new version
+                if challenge.has_data and not self.agent_ranks and not is_new_ver:
                     agents = await client.get_zzz_agents(self.account.uid)
                     self.agent_ranks = {agent.id: agent.rank for agent in agents}
-                challenge = await client.get_shiyu_defense(self.account.uid, previous=previous)
             elif self.challenge_type is ChallengeType.ASSAULT:
-                challenge = await client.get_deadly_assault(self.account.uid, previous=previous)
+                raw = await client.get_deadly_assault(self.account.uid, previous=previous, raw=True)
             else:
-                msg = f"Invalid challenge type: {self.challenge_type}"
-                raise ValueError(msg)
+                msg = f"Fetching data for {self.challenge_type!r}"
+                raise NotImplementedError(msg)
+
+            challenge = ChallengeHistory.load_data(raw, challenge_type=self.challenge_type)
+
+            if (
+                self.challenge_type in {ChallengeType.SPIRAL_ABYSS, ChallengeType.IMG_THEATER}
+                and not self.characters
+            ):
+                # Only fetch characters when challenge has data
+                try:
+                    self.check_challenge_data(challenge)
+                except NoChallengeDataError:
+                    pass
+                else:
+                    self.characters = await client.get_genshin_characters(self.account.uid)
 
             try:
                 season_id = self._get_season_id(challenge, previous)
@@ -276,18 +304,20 @@ async def _fetch_data(self) -> None:
                 uid=self.account.uid,
                 challenge_type=self.challenge_type,
                 season_id=season_id,
-                data=challenge,
+                raw=raw,
                 lang=client.lang,
             )
 
     def check_challenge_data(self, challenge: Challenge | None) -> None:
+        """Check if the challenge has data and raise an error if it doesn't"""
+        exc = NoChallengeDataError(self.challenge_type)
         if challenge is None:
-            raise NoChallengeDataError(self.challenge_type)
+            raise exc
         if isinstance(challenge, SpiralAbyss):
             if not challenge.floors:
                 raise NoChallengeDataError(ChallengeType.SPIRAL_ABYSS)
         elif not challenge.has_data:
-            raise NoChallengeDataError(self.challenge_type)
+            raise exc
 
     def get_season(self, challenge: Challenge) -> StarRailChallengeSeason:
         if isinstance(challenge, SpiralAbyss | ImgTheaterData | ShiyuDefense | DeadlyAssault):

migrations/models/33_20250325225453_update.py

@@ -0,0 +1,13 @@
+from tortoise import BaseDBAsyncClient
+
+
+async def upgrade(db: BaseDBAsyncClient) -> str:
+    return """
+        ALTER TABLE "challengehistory" ADD "json_data" JSONB;
+        ALTER TABLE "challengehistory" ALTER COLUMN "data" DROP NOT NULL;"""
+
+
+async def downgrade(db: BaseDBAsyncClient) -> str:
+    return """
+        ALTER TABLE "challengehistory" DROP COLUMN "json_data";
+        ALTER TABLE "challengehistory" ALTER COLUMN "data" SET NOT NULL;"""