feat(schedule): Add format flag

2 files changed, 15 insertions, 2 deletions

diff --git a/src/oguri/cli.py b/src/oguri/cli.py
index 6a624e9..e496dd1 100644
--- a/src/oguri/cli.py
+++ b/src/oguri/cli.py
@@ -34,7 +34,17 @@ def cli():
     help="Filter by country of origin. Can be used multiple times. Not an exhaustive list.",
 )
 @click.option("--episode", type=int, help="Filter by episode number.")
-def schedule(day, reverse, exact, first, sort_by, countries, episode):
[email protected](
+    "--format",
+    "formats",
+    multiple=True,
+    type=click.Choice(
+        ["TV", "TV_SHORT", "MOVIE", "SPECIAL", "OVA"], case_sensitive=False
+    ),
+    default=["TV", "TV_SHORT", "MOVIE", "SPECIAL", "OVA"],
+    help="Filter by format. Can be used multiple times.",
+)
+def schedule(day, reverse, exact, first, sort_by, countries, episode, formats):
     """
     Shows the airing schedule for a given day.
 
@@ -57,7 +67,7 @@ def schedule(day, reverse, exact, first, sort_by, countries, episode):
 
     asyncio.run(
         schedule_logic.show_schedule(
-            days_offset, reverse, exact, first, sort_by, countries, episode
+            days_offset, reverse, exact, first, sort_by, countries, episode, formats
         )
     )
 
diff --git a/src/oguri/schedule.py b/src/oguri/schedule.py
index 139430f..d16a4b0 100644
--- a/src/oguri/schedule.py
+++ b/src/oguri/schedule.py
@@ -19,6 +19,7 @@ async def show_schedule(
     sort_by: str = "time",
     countries: list[str] = ["JP", "KR", "CN"],
     episode: int | None = None,
+    formats: list[str] = ["TV", "TV_SHORT", "MOVIE", "SPECIAL", "OVA"],
 ):
     client = Client(url="https://graphql.anilist.co")
     start_of_day = datetime.now().replace(
@@ -41,6 +42,7 @@ async def show_schedule(
                     MediaTitleFields.native(),
                 ),
                 MediaFields.country_of_origin,
+                MediaFields.format,
             ),
         )
     )
@@ -61,6 +63,7 @@ async def show_schedule(
                         s
                         for s in airing_schedules
                         if s.get("media", {}).get("countryOfOrigin") in countries
+                        and s.get("media", {}).get("format") in formats
                     ]
 
                     if sort_by == "time":