Reformat. Add check for admin role for admin commands

BrigadaSOS · Jan 27, 2024 · 47b8965 · 47b8965
1 parent 492e981
commit 47b8965
Show file tree

Hide file tree

Showing 11 changed files with 114 additions and 110 deletions.
diff --git a/.env.example b/.env.example
@@ -1,5 +1,6 @@
 TOKEN=
 GUILD_ID=
 CHANNEL_ID=
+ADMIN_ROLE_ID=
 PROD_DB_PATH=./db/prod.db
 GOALS_DB_PATH=./db/goals.db
diff --git a/db/prod.db b/db/prod.db
diff --git a/immersion_bot/cogs/cogs_manager.py b/immersion_bot/cogs/cogs_manager.py
@@ -7,14 +7,15 @@
 from discord.ui import Select, View
 
 GUILD_ID = int(os.environ["GUILD_ID"])
+ADMIN_ROLE_ID = os.environ["ADMIN_ROLE_ID"]
 
 
 class BotManager(commands.Cog):
     def __init__(self, bot: commands.Bot) -> None:
         self.bot = bot
 
     @app_commands.command(name="load", description="Loads cogs.")
-    @app_commands.checks.has_permissions(administrator=True)
+    @app_commands.checks.has_role(ADMIN_ROLE_ID)
     async def load(self, interaction: discord.Interaction, *, cog: str):
         await interaction.response.defer()
         try:
@@ -27,7 +28,7 @@ async def load(self, interaction: discord.Interaction, *, cog: str):
             await interaction.edit_original_response(content="**`SUCCESS`**")
 
     @app_commands.command(name="stop", description="Stops cogs.")
-    @app_commands.checks.has_permissions(administrator=True)
+    @app_commands.checks.has_role(ADMIN_ROLE_ID)
     async def stop(self, interaction: discord.Interaction):
         if interaction.command_failed:
             await interaction.response.send_message(
@@ -58,22 +59,22 @@ async def my_callback(interaction):
     @app_commands.command(
         name="sync_global_commands", description="Syncs global slash commands."
     )
-    @app_commands.checks.has_permissions(administrator=True)
+    @app_commands.checks.has_role(ADMIN_ROLE_ID)
     async def sync_global_commands(self, interaction: discord.Interaction):
         await self.bot.tree.sync()
         await interaction.response.send_message(f"Synced global commands")
 
     @app_commands.command(
         name="clear_global_commands", description="Clears all global commands."
     )
-    @app_commands.checks.has_permissions(administrator=True)
+    @app_commands.checks.has_role(ADMIN_ROLE_ID)
     async def clear_global_commands(self, interaction: discord.Interaction):
         self.bot.tree.clear_commands(guild=None)
         await self.bot.tree.sync()
         await interaction.response.send_message("Cleared global commands.")
 
     @app_commands.command(name="sync", description="Syncs slash commands to the guild.")
-    @app_commands.checks.has_permissions(administrator=True)
+    @app_commands.checks.has_role(ADMIN_ROLE_ID)
     async def sync(self, interaction: discord.Interaction):
         self.bot.tree.copy_global_to(guild=discord.Object(id=GUILD_ID))
         await self.bot.tree.sync(guild=discord.Object(id=GUILD_ID))
@@ -84,7 +85,7 @@ async def sync(self, interaction: discord.Interaction):
     @app_commands.command(
         name="clear_guild_commands", description="Clears all guild commands."
     )
-    @app_commands.checks.has_permissions(administrator=True)
+    @app_commands.checks.has_role(ADMIN_ROLE_ID)
     async def clear_guild_commands(self, interaction: discord.Interaction):
         self.bot.tree.clear_commands(guild=discord.Object(id=GUILD_ID))
         await self.bot.tree.sync(guild=discord.Object(id=GUILD_ID))

diff --git a/immersion_bot/cogs/export.py b/immersion_bot/cogs/export.py
@@ -51,9 +51,7 @@ async def on_ready(self):
             Choice(name="Yearly", value="Yearly"),
         ]
     )
-    @app_commands.choices(
-        media_type=helpers.get_logeable_media_type_choices()
-    )
+    @app_commands.choices(media_type=helpers.get_logeable_media_type_choices())
     @app_commands.describe(
         date="""See past user overviews, combine it wit timeframes: [year-month-day] Example: '2022-12-29'."""
     )
@@ -96,7 +94,8 @@ async def export(
         for i, row in enumerate(logs):
             worksheet.write("A" + str(row_Index), row.media_type.value)
             worksheet.write(
-                "B" + str(row_Index), helpers._to_amount(row.media_type.value, row.amount)
+                "B" + str(row_Index),
+                helpers._to_amount(row.media_type.value, row.amount),
             )
             worksheet.write("C" + str(row_Index), str(row.note))
             worksheet.write("D" + str(row_Index), str(row.created_at))

diff --git a/immersion_bot/cogs/goal.py b/immersion_bot/cogs/goal.py
@@ -105,9 +105,7 @@ async def on_ready(self):
     @app_commands.describe(
         amount="""Episode to watch, characters or pages to read. Time to read/listen in [hr:min] or [min] for example '1.30' or '25'."""
     )
-    @app_commands.choices(
-        media_type=helpers.get_logeable_media_type_choices()
-    )
+    @app_commands.choices(media_type=helpers.get_logeable_media_type_choices())
     @app_commands.describe(
         name="""You can use vndb IDs for VN and Anilist codes for Anime, Manga and Light Novels"""
     )
@@ -184,7 +182,10 @@ async def set_goal(
     )
     @app_commands.describe(amount="""Points to log.""")
     @app_commands.choices(
-        media_type=[*helpers.get_logeable_media_type_choices(), Choice(name="Anything", value=MediaType.ANYTHING.value)]
+        media_type=[
+            *helpers.get_logeable_media_type_choices(),
+            Choice(name="Anything", value=MediaType.ANYTHING.value),
+        ]
     )
     @app_commands.choices(frequency=[Choice(name="Daily", value="Daily")])
     @app_commands.describe(frequency="Make this your daily goal for the month.")
@@ -219,7 +220,7 @@ async def set_goal_points(
             f"de {media_type.upper()}",
             datetime.now(),
             frequency,
-            "points", # Do no translate, this is the table name
+            "points",  # Do no translate, this is the table name
         )
         if bool:
             return await interaction.response.send_message(

diff --git a/immersion_bot/cogs/leaderboard.py b/immersion_bot/cogs/leaderboard.py
@@ -40,9 +40,7 @@ async def on_ready(self):
             Choice(name="Yearly", value="Yearly"),
         ]
     )
-    @app_commands.choices(
-        media_type=helpers.get_logeable_media_type_choices()
-    )
+    @app_commands.choices(media_type=helpers.get_logeable_media_type_choices())
     @app_commands.describe(
         date="""See past leaderboards, combine it wit timeframes: [year-month-day] Example: '2022-12-29'."""
     )

diff --git a/immersion_bot/cogs/log.py b/immersion_bot/cogs/log.py
@@ -187,13 +187,13 @@ async def log_vn(
         backlog="""Registra en un día anterior. Formato aceptado: [yyyy-mm-dd]. Ejemplo: 2024-02-01"""
     )
     async def log_ln(
-            self,
-            interaction: discord.Interaction,
-            caracteres: int,
-            tiempo: str,
-            nombre: str,
-            comentario: Optional[str],
-            backlog: Optional[str],
+        self,
+        interaction: discord.Interaction,
+        caracteres: int,
+        tiempo: str,
+        nombre: str,
+        comentario: Optional[str],
+        backlog: Optional[str],
     ):
         await interaction.response.defer()
 
@@ -231,13 +231,13 @@ async def log_ln(
     @app_commands.describe(
         backlog="""Registra en un día anterior. Formato aceptado: [yyyy-mm-dd]. Ejemplo: 2024-02-01"""
     )
-    async def log_listening(
-            self,
-            interaction: discord.Interaction,
-            tiempo: str,
-            nombre: str,
-            comentario: Optional[str],
-            backlog: Optional[str],
+    async def log_readtime(
+        self,
+        interaction: discord.Interaction,
+        tiempo: str,
+        nombre: str,
+        comentario: Optional[str],
+        backlog: Optional[str],
     ):
         await interaction.response.defer()
 
@@ -538,72 +538,75 @@ async def log_autocomplete(
         interaction: discord.Interaction,
         current: str,
     ) -> List[app_commands.Choice[str]]:
-        await interaction.response.defer()
+        print("Autocompleting")
         media_type = interaction.namespace["media_type"]
         suggestions = []
         url = ""
 
-        if media_type == "VN":
-            url = "https://api.vndb.org/kana/vn"
-            data = {
-                "filters": ["search", "=", f"{current}"],
-                "fields": "title, alttitle",
-            }  # default no. of results is 10
-
-        elif media_type == "Anime" or media_type == "Manga":
-            url = "https://graphql.anilist.co"
-            query = f"""
-            query ($page: Int, $perPage: Int, $title: String) {{
-                Page(page: $page, perPage: $perPage) {{
-                    pageInfo {{
-                        total
-                        perPage
-                    }}
-                    media (search: $title, type: {media_type.upper()}) {{
-                        id
-                        title {{
-                            romaji
-                            native
-                        }}
-                    }}
-                }}
-            }}
-            """
-
-            variables = {"title": current, "page": 1, "perPage": 10}
-
-            data = {"query": query, "variables": variables}
-
-        if not url:
-            return []
-
-        async with aiohttp.ClientSession() as session:
-            async with session.post(url, json=data) as resp:
-                log.info(resp.status)
-                json_data = await resp.json()
-
-                if media_type == "VN":
-                    suggestions = [
-                        (result["title"], result["id"])
-                        for result in json_data["results"]
-                    ]
-
-                elif media_type == "Anime" or media_type == "Manga":
-                    suggestions = [
-                        (
-                            f"{result['title']['romaji']} ({result['title']['native']})",
-                            result["id"],
-                        )
-                        for result in json_data["data"]["Page"]["media"]
-                    ]
-
-                await asyncio.sleep(0)
-
-                return [
-                    app_commands.Choice(name=title, value=str(id))
-                    for title, id in suggestions
-                    if current.lower() in title.lower()
-                ]
+        return []
+        #
+        # if media_type == "VN":
+        #     url = "https://api.vndb.org/kana/vn"
+        #     data = {
+        #         "filters": ["search", "=", f"{current}"],
+        #         "fields": "title, alttitle",
+        #     }  # default no. of results is 10
+        #
+        # elif media_type == "Anime" or media_type == "Manga":
+        #     url = "https://graphql.anilist.co"
+        #     query = f"""
+        #     query ($page: Int, $perPage: Int, $title: String) {{
+        #         Page(page: $page, perPage: $perPage) {{
+        #             pageInfo {{
+        #                 total
+        #                 perPage
+        #             }}
+        #             media (search: $title, type: {media_type.upper()}) {{
+        #                 id
+        #                 title {{
+        #                     romaji
+        #                     native
+        #                 }}
+        #             }}
+        #         }}
+        #     }}
+        #     """
+        #
+        #     variables = {"title": current, "page": 1, "perPage": 10}
+        #
+        #     data = {"query": query, "variables": variables}
+        #
+        # if not url:
+        #     return []
+        #
+        # async with aiohttp.ClientSession() as session:
+        #     async with session.post(url, json=data) as resp:
+        #         log.info(resp.status)
+        #         json_data = await resp.json()
+        #
+        #         if media_type == "VN":
+        #             suggestions = [
+        #                 (result["title"], result["id"])
+        #                 for result in json_data["results"]
+        #             ]
+        #
+        #         elif media_type == "Anime" or media_type == "Manga":
+        #             suggestions = [
+        #                 (
+        #                     f"{result['title']['romaji']} ({result['title']['native']})",
+        #                     result["id"],
+        #                 )
+        #                 for result in json_data["data"]["Page"]["media"]
+        #             ]
+        #
+        #         await asyncio.sleep(0)
+        #
+        #         return [
+        #             app_commands.Choice(name=title, value=str(id))
+        #             for title, id in suggestions
+        #             if current.lower() in title.lower()
+        #         ]
+        #
 
 
 async def setup(bot: commands.Bot) -> None:

diff --git a/immersion_bot/cogs/user.py b/immersion_bot/cogs/user.py
@@ -42,6 +42,7 @@ def __conform__(self, protocol):
         if protocol is sqlite3.PrepareProtocol:
             return self.name
 
+
 class User(commands.Cog):
     def __init__(self, bot: commands.Bot) -> None:
         self.bot = bot
@@ -193,9 +194,7 @@ async def create_embed(self, timeframe, interaction, weighed_points_mediums, log
             Choice(name="Año", value="Yearly"),
         ]
     )
-    @app_commands.choices(
-        media_type=helpers.get_logeable_media_type_choices()
-    )
+    @app_commands.choices(media_type=helpers.get_logeable_media_type_choices())
     @app_commands.describe(
         date="""See past user overviews, combine it wit timeframes: [year-month-day] Example: '2022-12-29'."""
     )

diff --git a/immersion_bot/helpers.py b/immersion_bot/helpers.py
@@ -20,6 +20,7 @@ def __conform__(self, protocol):
         if protocol is sqlite3.PrepareProtocol:
             return self.name
 
+
 def _to_amount(media_type, amount):
     if media_type == MediaType.ANIME.value:
         return amount * 9.5
@@ -38,14 +39,15 @@ def _to_amount(media_type, amount):
 
 
 def get_logeable_media_type_choices():
-      return [
-            Choice(name="Anime", value=MediaType.ANIME.value),
-            Choice(name="Manga", value=MediaType.MANGA.value),
-            Choice(name="Visual Novels", value=MediaType.VN.value),
-            Choice(name="Light Novels", value=MediaType.LN.value),
-            Choice(name="Readtime", value=MediaType.READTIME.value),
-            Choice(name="Listening", value=MediaType.LISTENING.value),
-        ]
+    return [
+        Choice(name="Anime", value=MediaType.ANIME.value),
+        Choice(name="Manga", value=MediaType.MANGA.value),
+        Choice(name="Visual Novels", value=MediaType.VN.value),
+        Choice(name="Light Novels", value=MediaType.LN.value),
+        Choice(name="Readtime", value=MediaType.READTIME.value),
+        Choice(name="Listening", value=MediaType.LISTENING.value),
+    ]
+
 
 def multiplied_points(logs):
     dictes = defaultdict(list)

diff --git a/immersion_bot/launch_bot.py b/immersion_bot/launch_bot.py
@@ -47,7 +47,7 @@ def __init__(self) -> None:
             intents=discord.Intents.all(),
             tree_cls=CustomCommandTree,
         )
-        locale.setlocale(locale.LC_TIME, 'es_ES.UTF-8')
+        locale.setlocale(locale.LC_TIME, "es_ES.UTF-8")
 
     async def on_error(self, event_method: str, /, *args, **kwargs):
         log.exception("Ignoring exception in %s", event_method)

diff --git a/immersion_bot/sql.py b/immersion_bot/sql.py
@@ -11,11 +11,11 @@ def __conform__(self, protocol):
 
 class MediaType(SqliteEnum):
     # As it is
-    ANIME = "ANIME" # Done
-    MANGA = "MANGA" # Done
-    VN = "VN" # Done
+    ANIME = "ANIME"  # Done
+    MANGA = "MANGA"  # Done
+    VN = "VN"  # Done
     LN = "LN"
-    LISTENING = "LISTENING" # Done
+    LISTENING = "LISTENING"  # Done
     READTIME = "READTIME"
     ANYTHING = "ANYTHING"  # ANYTHING as an option for setting a point goal