feat(schedule): Support pagination

1 files changed, 104 insertions, 85 deletions

diff --git a/src/oguri/schedule.py b/src/oguri/schedule.py
index 1d48a54..c56db98 100644
--- a/src/oguri/schedule.py
+++ b/src/oguri/schedule.py
@@ -7,6 +7,7 @@ from anilist_client.custom_fields import (
     MediaFields,
     MediaTitleFields,
     PageFields,
+    PageInfoFields,
 )
 from anilist_client.custom_queries import Query
 from . import constants
@@ -46,91 +47,109 @@ async def show_schedule(
     end_of_day = start_of_day.replace(
         hour=23, minute=59, second=59, microsecond=999999)
 
-  airing_schedules_query = Query.page().fields(
-      PageFields.airing_schedules(
-          airing_at_greater=int(start_of_day.timestamp()),
-          episode=episode if episode else 1 if first_episode_only else None,
-          airing_at_lesser=int(end_of_day.timestamp()),
-      ).fields(
-          AiringScheduleFields.airing_at,
-          AiringScheduleFields.episode,
-          AiringScheduleFields.media().fields(
-              MediaFields.site_url,
-              MediaFields.title().fields(
-                  MediaTitleFields.english(),
-                  MediaTitleFields.romaji(),
-                  MediaTitleFields.native(),
-              ),
-              MediaFields.country_of_origin,
-              MediaFields.format,
-              MediaFields.source(),
-              MediaFields.genres,
-          ),
-      ))
-
-  try:
-    response = await client.query(
-        airing_schedules_query, operation_name="get_airing_schedules")
-
-    if response:
-      page = response.get("Page")
-
-      if page:
-        airing_schedules = page.get("airingSchedules")
-
-        if airing_schedules:
-          airing_schedules = [
-              s for s in airing_schedules
-              if s.get("media", {}).get("countryOfOrigin") in countries and
-              s.get("media", {}).get("format") in formats and
-              s.get("media", {}).get("source") in sources and any(
-                  g in genres for g in s.get("media", {}).get("genres", []))
-          ]
-
-          if sort_by == "time":
-            airing_schedules.sort(
-                key=lambda x: x.get("airingAt"), reverse=not reverse_order)
-          elif sort_by == "episode":
-            airing_schedules.sort(
-                key=lambda x: x.get("episode"), reverse=not reverse_order)
-
-          console = Console()
-          table = Table(show_header=True, header_style="bold magenta")
-
-          table.add_column("Title")
-          table.add_column("Episode")
-          table.add_column("Airing Time")
-
-          for schedule in airing_schedules:
-            airing_at = schedule.get("airingAt")
-            episode = schedule.get("episode")
-            media = schedule.get("media")
-            site_url = media.get("siteUrl")
-            titles = media.get("title")
-            title = (
-                titles.get("english") or titles.get("romaji") or
-                titles.get("native"))
-
-            if airing_at:
-              airing_at_date = datetime.fromtimestamp(airing_at)
-              airing_time_string = ""
-
-              if exact_time:
-                airing_time_string = format_exact_time(airing_at_date)
-              elif datetime.now() > airing_at_date:
-                airing_time_string = (f"Aired {relative_time(airing_at_date)}")
-              else:
-                airing_time_string = format_future_airing_time(airing_at_date)
-
-              table.add_row(
-                  f"[link={site_url}]{title}[/link]",
-                  str(episode),
-                  airing_time_string,
-              )
-
-          console.print(table)
-  except Exception as exception:
-    print(exception)
+  all_airing_schedules = []
+  has_next_page = True
+  page_number = 1
+
+  while has_next_page:
+    airing_schedules_query = Query.page(page=page_number).fields(
+        PageFields.airing_schedules(
+            airing_at_greater=int(start_of_day.timestamp()),
+            episode=episode if episode else 1 if first_episode_only else None,
+            airing_at_lesser=int(end_of_day.timestamp()),
+        ).fields(
+            AiringScheduleFields.airing_at,
+            AiringScheduleFields.episode,
+            AiringScheduleFields.media().fields(
+                MediaFields.site_url,
+                MediaFields.title().fields(
+                    MediaTitleFields.english(),
+                    MediaTitleFields.romaji(),
+                    MediaTitleFields.native(),
+                ),
+                MediaFields.country_of_origin,
+                MediaFields.format,
+                MediaFields.source(),
+                MediaFields.genres,
+            ),
+        ),
+        PageFields.page_info().fields(PageInfoFields.has_next_page))
+
+    try:
+      response = await client.query(
+          airing_schedules_query, operation_name="get_airing_schedules")
+
+      if response:
+        page = response.get("Page")
+
+        if page:
+          airing_schedules = page.get("airingSchedules")
+          page_info = page.get("pageInfo")
+          has_next_page = page_info.get("hasNextPage") if page_info else False
+
+          if airing_schedules:
+            all_airing_schedules.extend(airing_schedules)
+
+          page_number += 1
+        else:
+          has_next_page = False
+      else:
+        has_next_page = False
+    except Exception as exception:
+      print(exception)
+
+      has_next_page = False
+
+  if all_airing_schedules:
+    airing_schedules = [
+        s for s in all_airing_schedules
+        if s.get("media", {}).get("countryOfOrigin") in countries and
+        s.get("media", {}).get("format") in formats and
+        s.get("media", {}).get("source") in sources and any(
+            g in genres for g in s.get("media", {}).get("genres", []))
+    ]
+
+    if sort_by == "time":
+      airing_schedules.sort(
+          key=lambda x: x.get("airingAt"), reverse=not reverse_order)
+    elif sort_by == "episode":
+      airing_schedules.sort(
+          key=lambda x: x.get("episode"), reverse=not reverse_order)
+
+    console = Console()
+    table = Table(show_header=True, header_style="bold magenta")
+
+    table.add_column("Title")
+    table.add_column("Episode")
+    table.add_column("Airing Time")
+
+    for schedule in airing_schedules:
+      airing_at = schedule.get("airingAt")
+      episode = schedule.get("episode")
+      media = schedule.get("media")
+      site_url = media.get("siteUrl")
+      titles = media.get("title")
+      title = (
+          titles.get("english") or titles.get("romaji") or titles.get("native"))
+
+      if airing_at:
+        airing_at_date = datetime.fromtimestamp(airing_at)
+        airing_time_string = ""
+
+        if exact_time:
+          airing_time_string = format_exact_time(airing_at_date)
+        elif datetime.now() > airing_at_date:
+          airing_time_string = (f"Aired {relative_time(airing_at_date)}")
+        else:
+          airing_time_string = format_future_airing_time(airing_at_date)
+
+        table.add_row(
+            f"[link={site_url}]{title}[/link]",
+            str(episode),
+            airing_time_string,
+        )
+
+    console.print(table)
 
 
 def format_exact_time(date):