enable querying world's end songs

cogs/search.py

@@ -1,3 +1,4 @@
+import shlex
 from datetime import datetime
 from typing import TYPE_CHECKING
 
@@ -11,6 +12,7 @@
 from chunithm_net.consts import JACKET_BASE
 from database.models import Alias, Chart, Song
 from utils import (
+    Arguments,
     did_you_mean_text,
     release_to_chunithm_version,
     yt_search_link,
@@ -155,11 +157,21 @@ async def info(self, ctx: Context, *, query: str):
         query: str
             The song title or alias to search for.
         """
+        parser = Arguments()
+        parser.add_argument("query", nargs="+")
+        parser.add_argument("-we", "--worlds-end", action="store_true")
+
+        try:
+            args = parser.parse_intermixed_args(shlex.split(query))
+            query = " ".join(args.query)
+        except RuntimeError as e:
+            await ctx.reply(str(e), mention_author=False)
+            return
 
         async with ctx.typing(), self.bot.begin_db_session() as session:
             guild_id = ctx.guild.id if ctx.guild is not None else None
             song, alias, similarity = await self.utils.find_song(
-                query, guild_id=guild_id
+                query, guild_id=guild_id, worlds_end=args.worlds_end
             )
 
             if similarity < 0.9:
@@ -197,8 +209,9 @@ async def info(self, ctx: Context, *, query: str):
                     if chart.sdvxin_chart_view is not None
                     else yt_search_link(song.title, chart.difficulty)
                 )
-                desc = f"[{chart.difficulty[0]}]({url}) {chart.level}"
-                if chart.const != 0:
+                worlds_end = "WORLD'S END"
+                desc = f'[{worlds_end if args.worlds_end else chart.difficulty[0]}]({url}) {chart.level}'
+                if chart.const is not None:
                     desc += f" ({chart.const:.1f})"
                 chart_level_desc.append(desc)
 

utils/__init__.py

@@ -36,6 +36,8 @@ def did_you_mean_text(result: "Song", alias: "Alias | None") -> str:
 
 
 def yt_search_link(title: str, difficulty: str) -> str:
+    if difficulty == "WE":
+        difficulty = "WORLD'S END"
     return "https://www.youtube.com/results?search_query=" + quote(
         f"CHUNITHM {title} {difficulty}"
     )