feat(schedule): Add country flag

2 files changed, 20 insertions, 2 deletions

diff --git a/src/oguri/cli.py b/src/oguri/cli.py
index c3327f6..dc3fbb6 100644
--- a/src/oguri/cli.py
+++ b/src/oguri/cli.py
@@ -25,7 +25,15 @@ def cli():
     default="time",
     help="Sort the schedule by time or episode.",
 )
-def schedule(day, reverse, exact, first, sort_by):
[email protected](
+    "--country",
+    "countries",
+    multiple=True,
+    type=click.Choice(["JP", "KR", "CN"], case_sensitive=False),
+    default=["JP", "KR", "CN"],
+    help="Filter by country of origin. Can be used multiple times. Not an exhaustive list.",
+)
+def schedule(day, reverse, exact, first, sort_by, countries):
     """
     Shows the airing schedule for a given day.
 
@@ -47,7 +55,9 @@ def schedule(day, reverse, exact, first, sort_by):
             )
 
     asyncio.run(
-        schedule_logic.show_schedule(days_offset, reverse, exact, first, sort_by)
+        schedule_logic.show_schedule(
+            days_offset, reverse, exact, first, sort_by, countries
+        )
     )
 
 
diff --git a/src/oguri/schedule.py b/src/oguri/schedule.py
index c6f837a..07a48e8 100644
--- a/src/oguri/schedule.py
+++ b/src/oguri/schedule.py
@@ -17,6 +17,7 @@ async def show_schedule(
     exact_time: bool = False,
     first_episode_only: bool = False,
     sort_by: str = "time",
+    countries: list[str] = ["JP", "KR", "CN"],
 ):
     client = Client(url="https://graphql.anilist.co")
     start_of_day = datetime.now().replace(
@@ -38,6 +39,7 @@ async def show_schedule(
                     MediaTitleFields.romaji(),
                     MediaTitleFields.native(),
                 ),
+                MediaFields.country_of_origin,
             ),
         )
     )
@@ -54,6 +56,12 @@ async def show_schedule(
                 airing_schedules = page.get("airingSchedules")
 
                 if airing_schedules:
+                    airing_schedules = [
+                        s
+                        for s in airing_schedules
+                        if s.get("media", {}).get("countryOfOrigin") in countries
+                    ]
+
                     if sort_by == "time":
                         airing_schedules.sort(
                             key=lambda x: x.get("airingAt"), reverse=not reverse_order