feat(schedule): Add genre flag

2 files changed, 110 insertions, 11 deletions

diff --git a/src/oguri/cli.py b/src/oguri/cli.py
index d34dbb3..b8aa41a 100644
--- a/src/oguri/cli.py
+++ b/src/oguri/cli.py
@@ -57,7 +57,71 @@ def cli():
 )
 @click.option("--year", type=int, help="Filter by airing year.")
 @click.option("--month", type=int, help="Filter by airing month (1-12).")
-def schedule(day, reverse, exact, first, sort_by, countries, episode, formats, sources, year, month):
[email protected](
+    "--genre",
+    "genres",
+    multiple=True,
+    type=click.Choice(
+        [
+            "Action",
+            "Adventure",
+            "Comedy",
+            "Drama",
+            "Ecchi",
+            "Fantasy",
+            "Hentai",
+            "Horror",
+            "Mahou Shoujo",
+            "Mecha",
+            "Music",
+            "Mystery",
+            "Psychological",
+            "Romance",
+            "Sci-Fi",
+            "Slice of Life",
+            "Sports",
+            "Supernatural",
+            "Thriller",
+        ],
+        case_sensitive=False,
+    ),
+    default=[
+        "Action",
+        "Adventure",
+        "Comedy",
+        "Drama",
+        "Ecchi",
+        "Fantasy",
+        "Hentai",
+        "Horror",
+        "Mahou Shoujo",
+        "Mecha",
+        "Music",
+        "Mystery",
+        "Psychological",
+        "Romance",
+        "Sci-Fi",
+        "Slice of Life",
+        "Sports",
+        "Supernatural",
+        "Thriller",
+    ],
+    help="Filter by genre. Can be used multiple times.",
+)
+def schedule(
+    day,
+    reverse,
+    exact,
+    first,
+    sort_by,
+    countries,
+    episode,
+    formats,
+    sources,
+    year,
+    month,
+    genres,
+):
     """
     Shows the airing schedule for a given day.
 
@@ -91,6 +155,7 @@ def schedule(day, reverse, exact, first, sort_by, countries, episode, formats, s
             sources,
             year,
             month,
+            genres,
         )
     )
 
diff --git a/src/oguri/schedule.py b/src/oguri/schedule.py
index ba3be64..35819cd 100644
--- a/src/oguri/schedule.py
+++ b/src/oguri/schedule.py
@@ -30,22 +30,52 @@ async def show_schedule(
     ],
     year: int | None = None,
     month: int | None = None,
+    genres: list[str] = [
+        "Action",
+        "Adventure",
+        "Comedy",
+        "Drama",
+        "Ecchi",
+        "Fantasy",
+        "Hentai",
+        "Horror",
+        "Mahou Shoujo",
+        "Mecha",
+        "Music",
+        "Mystery",
+        "Psychological",
+        "Romance",
+        "Sci-Fi",
+        "Slice of Life",
+        "Sports",
+        "Supernatural",
+        "Thriller",
+    ],
 ):
     client = Client(url="https://graphql.anilist.co")
-    if year and month:
-        start_of_day = datetime(year, month, 1).replace(
+    now = datetime.now()
+    target_year = year if year is not None else now.year
+    target_month = month if month is not None else now.month
+
+    if days_offset != 0:
+        start_of_day = now.replace(
             hour=0, minute=0, second=0, microsecond=0
+        ) + timedelta(days=days_offset)
+        end_of_day = start_of_day.replace(
+            hour=23, minute=59, second=59, microsecond=999999
         )
-        # Calculate the last day of the month
-        if month == 12:
-            end_of_day = datetime(year + 1, 1, 1) - timedelta(microseconds=1)
-        else:
-            end_of_day = datetime(year, month + 1, 1) - timedelta(microseconds=1)
     else:
-        start_of_day = datetime.now().replace(
+        start_of_day = datetime(target_year, target_month, 1).replace(
             hour=0, minute=0, second=0, microsecond=0
-        ) + timedelta(days=days_offset)
-        end_of_day = start_of_day.replace(hour=23, minute=59, second=59, microsecond=999999)
+        )
+
+        if target_month == 12:
+            end_of_day = datetime(target_year + 1, 1, 1) - timedelta(microseconds=1)
+        else:
+            end_of_day = datetime(target_year, target_month + 1, 1) - timedelta(
+                microseconds=1
+            )
+
     airing_schedules_query = Query.page().fields(
         PageFields.airing_schedules(
             airing_at_greater=int(start_of_day.timestamp()),
@@ -64,6 +94,7 @@ async def show_schedule(
                 MediaFields.country_of_origin,
                 MediaFields.format,
                 MediaFields.source(),
+                MediaFields.genres,
             ),
         )
     )
@@ -86,6 +117,9 @@ async def show_schedule(
                         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":