Allow returning raw data for notes

thesadru · Nov 22, 2024 · 1202cc9 · 1202cc9
1 parent ec5dba9
commit 1202cc9
Show file tree

Hide file tree

Showing 4 changed files with 88 additions and 4 deletions.
diff --git a/genshin/client/components/chronicle/genshin.py b/genshin/client/components/chronicle/genshin.py
@@ -169,13 +169,32 @@ async def get_imaginarium_theater(
 
         return models.ImgTheater(**data)
 
+    @typing.overload
+    async def get_genshin_notes(
+        self,
+        uid: typing.Optional[int] = ...,
+        *,
+        lang: typing.Optional[str] = ...,
+        autoauth: bool = ...,
+        return_raw_data: typing.Literal[False] = ...,
+    ) -> models.Notes: ...
+    @typing.overload
+    async def get_genshin_notes(
+        self,
+        uid: typing.Optional[int] = ...,
+        *,
+        lang: typing.Optional[str] = ...,
+        autoauth: bool = ...,
+        return_raw_data: typing.Literal[True] = ...,
+    ) -> typing.Mapping[str, typing.Any]: ...
     async def get_genshin_notes(
         self,
         uid: typing.Optional[int] = None,
         *,
         lang: typing.Optional[str] = None,
         autoauth: bool = True,
-    ) -> models.Notes:
+        return_raw_data: bool = False,
+    ) -> typing.Union[models.Notes, typing.Mapping[str, typing.Any]]:
         """Get genshin real-time notes."""
         try:
             data = await self._request_genshin_record("dailyNote", uid, lang=lang)
@@ -189,6 +208,8 @@ async def get_genshin_notes(
             await self.update_settings(3, True, game=types.Game.GENSHIN)
             data = await self._request_genshin_record("dailyNote", uid, lang=lang)
 
+        if return_raw_data:
+            return data
         return models.Notes(**data)
 
     async def get_genshin_activities(self, uid: int, *, lang: typing.Optional[str] = None) -> models.Activities:

diff --git a/genshin/client/components/chronicle/honkai.py b/genshin/client/components/chronicle/honkai.py
@@ -130,14 +130,33 @@ async def get_honkai_memorial_arena(
         data = await self._request_honkai_record("battleFieldReport", uid, lang=lang)
         return [models.MemorialArena(**x) for x in data["reports"]]
 
+    @typing.overload
+    async def get_honkai_notes(
+        self,
+        uid: int,
+        *,
+        lang: typing.Optional[str] = ...,
+        return_raw_data: typing.Literal[False] = ...
+    ) -> models.HonkaiNotes: ...
+    @typing.overload
+    async def get_honkai_notes(
+        self,
+        uid: int,
+        *,
+        lang: typing.Optional[str] = ...,
+        return_raw_data: typing.Literal[True] = ...
+    ) -> typing.Mapping[str, typing.Any]: ...
     async def get_honkai_notes(
         self,
         uid: int,
         *,
         lang: typing.Optional[str] = None,
-    ) -> models.HonkaiNotes:
+        return_raw_data: bool = False,
+    ) -> typing.Union[models.HonkaiNotes, typing.Mapping[str, typing.Any]]:
         """Get honkai memorial arena."""
         data = await self._request_honkai_record("note", uid, lang=lang)
+        if return_raw_data:
+            return data
         return models.HonkaiNotes(**data)
 
     async def get_full_honkai_user(

diff --git a/genshin/client/components/chronicle/starrail.py b/genshin/client/components/chronicle/starrail.py
@@ -57,13 +57,34 @@ async def _request_starrail_record(
             cache=cache_key,
         )
 
+    @typing.overload
+    async def get_starrail_notes(
+        self,
+        uid: typing.Optional[int] = ...,
+        *,
+        lang: typing.Optional[str] = ...,
+        autoauth: bool = ...,
+        return_raw_data: typing.Literal[False] = ...
+    ) -> models.StarRailNote:
+        ...
+    @typing.overload
+    async def get_starrail_notes(
+        self,
+        uid: typing.Optional[int] = ...,
+        *,
+        lang: typing.Optional[str] = ...,
+        autoauth: bool = ...,
+        return_raw_data: typing.Literal[True] = ...
+    ) -> typing.Mapping[str, typing.Any]:
+        ...
     async def get_starrail_notes(
         self,
         uid: typing.Optional[int] = None,
         *,
         lang: typing.Optional[str] = None,
         autoauth: bool = True,
-    ) -> models.StarRailNote:
+        return_raw_data: bool = False,
+    ) -> typing.Union[models.StarRailNote, typing.Mapping[str, typing.Any]]:
         """Get starrail real-time notes."""
         try:
             data = await self._request_starrail_record("note", uid, lang=lang)
@@ -77,6 +98,8 @@ async def get_starrail_notes(
             await self.update_settings(3, True, game=types.Game.STARRAIL)
             data = await self._request_starrail_record("note", uid, lang=lang)
 
+        if return_raw_data:
+            return data
         return models.StarRailNote(**data)
 
     async def get_starrail_user(

diff --git a/genshin/client/components/chronicle/zzz.py b/genshin/client/components/chronicle/zzz.py
@@ -71,13 +71,32 @@ async def _request_zzz_record(
             is_nap_ledger=is_nap_ledger,
         )
 
+    @typing.overload
+    async def get_zzz_notes(
+        self,
+        uid: typing.Optional[int] = ...,
+        *,
+        lang: typing.Optional[str] = ...,
+        autoauth: bool = ...,
+        return_raw_data: typing.Literal[False] = ...,
+    ) -> models.ZZZNotes: ...
+    @typing.overload
+    async def get_zzz_notes(
+        self,
+        uid: typing.Optional[int] = ...,
+        *,
+        lang: typing.Optional[str] = ...,
+        autoauth: bool = ...,
+        return_raw_data: typing.Literal[True] = ...,
+    ) -> typing.Mapping[str, typing.Any]: ...
     async def get_zzz_notes(
         self,
         uid: typing.Optional[int] = None,
         *,
         lang: typing.Optional[str] = None,
         autoauth: bool = True,
-    ) -> models.ZZZNotes:
+        return_raw_data: bool = False,
+    ) -> typing.Union[models.ZZZNotes, typing.Mapping[str, typing.Any]]:
         """Get ZZZ sticky notes (real-time notes)."""
         try:
             data = await self._request_zzz_record("note", uid, lang=lang)
@@ -91,6 +110,8 @@ async def get_zzz_notes(
             await self.update_settings(3, True, game=types.Game.ZZZ)
             data = await self._request_zzz_record("note", uid, lang=lang)
 
+        if return_raw_data:
+            return data
         return models.ZZZNotes(**data)
 
     async def get_zzz_diary(