From 8698e945bb7e767d44675c75c9edbb2bfc2a4e16 Mon Sep 17 00:00:00 2001 From: seria Date: Sat, 6 Jul 2024 06:37:29 +0900 Subject: [PATCH] Prevent enum crashes --- genshin/models/genshin/chronicle/notes.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/genshin/models/genshin/chronicle/notes.py b/genshin/models/genshin/chronicle/notes.py index aba31f80..0ab5b9d1 100644 --- a/genshin/models/genshin/chronicle/notes.py +++ b/genshin/models/genshin/chronicle/notes.py @@ -97,7 +97,14 @@ class TaskRewardStatus(str, enum.Enum): class TaskReward(APIModel): """Status of the Commission/Task.""" - status: TaskRewardStatus + status: typing.Union[TaskRewardStatus, str] + + @pydantic.validator("status", pre=True) + def __prevent_enum_crash(cls, v: str) -> typing.Union[TaskRewardStatus, str]: + try: + return TaskRewardStatus(v) + except ValueError: + return v class AttendanceRewardStatus(str, enum.Enum): @@ -112,9 +119,16 @@ class AttendanceRewardStatus(str, enum.Enum): class AttendanceReward(APIModel): """Status of the Encounter Point.""" - status: AttendanceRewardStatus + status: typing.Union[AttendanceRewardStatus, str] progress: int + @pydantic.validator("status", pre=True) + def __prevent_enum_crash(cls, v: str) -> typing.Union[AttendanceRewardStatus, str]: + try: + return AttendanceRewardStatus(v) + except ValueError: + return v + class DailyTasks(APIModel): """Daily tasks section."""