feat(schedule): Add source flag

2 files changed, 31 insertions, 2 deletions

diff --git a/src/oguri/cli.py b/src/oguri/cli.py
index e496dd1..6aa6cc9 100644
--- a/src/oguri/cli.py
+++ b/src/oguri/cli.py
@@ -44,7 +44,18 @@ def cli():
     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):
[email protected](
+    "--source",
+    "sources",
+    multiple=True,
+    type=click.Choice(
+        ["ORIGINAL", "MANGA", "LIGHT_NOVEL", "VISUAL_NOVEL", "VIDEO_GAME", "OTHER"],
+        case_sensitive=False,
+    ),
+    default=["ORIGINAL", "MANGA", "LIGHT_NOVEL", "VISUAL_NOVEL", "VIDEO_GAME", "OTHER"],
+    help="Filter by source. Can be used multiple times.",
+)
+def schedule(day, reverse, exact, first, sort_by, countries, episode, formats, sources):
     """
     Shows the airing schedule for a given day.
 
@@ -67,7 +78,15 @@ def schedule(day, reverse, exact, first, sort_by, countries, episode, formats):
 
     asyncio.run(
         schedule_logic.show_schedule(
-            days_offset, reverse, exact, first, sort_by, countries, episode, formats
+            days_offset,
+            reverse,
+            exact,
+            first,
+            sort_by,
+            countries,
+            episode,
+            formats,
+            sources,
         )
     )
 
diff --git a/src/oguri/schedule.py b/src/oguri/schedule.py
index d16a4b0..7a67df4 100644
--- a/src/oguri/schedule.py
+++ b/src/oguri/schedule.py
@@ -20,6 +20,14 @@ async def show_schedule(
     countries: list[str] = ["JP", "KR", "CN"],
     episode: int | None = None,
     formats: list[str] = ["TV", "TV_SHORT", "MOVIE", "SPECIAL", "OVA"],
+    sources: list[str] = [
+        "ORIGINAL",
+        "MANGA",
+        "LIGHT_NOVEL",
+        "VISUAL_NOVEL",
+        "VIDEO_GAME",
+        "OTHER",
+    ],
 ):
     client = Client(url="https://graphql.anilist.co")
     start_of_day = datetime.now().replace(
@@ -43,6 +51,7 @@ async def show_schedule(
                 ),
                 MediaFields.country_of_origin,
                 MediaFields.format,
+                MediaFields.source(),
             ),
         )
     )
@@ -64,6 +73,7 @@ async def show_schedule(
                         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
                     ]
 
                     if sort_by == "time":