feat(repo): make all cogs pylance-typechecking compliant

at `basic` level, does not include Aurora as it's being rewritten in the `aurora/v3` branch

pterodactyl/pterodactyl.py

@@ -1,6 +1,6 @@
 import asyncio
 import json
-from typing import Mapping, Optional, Tuple, Union
+from typing import AsyncIterable, Iterable, Mapping, Optional, Tuple, Union
 
 import discord
 import websockets
@@ -9,8 +9,9 @@ from pydactyl import PterodactylClient
 from redbot.core import app_commands, commands
 from redbot.core.app_commands import Choice
 from redbot.core.bot import Red
-from redbot.core.utils.chat_formatting import bold, box, error, humanize_list
+from redbot.core.utils.chat_formatting import bold, box, humanize_list
 from redbot.core.utils.views import ConfirmView
+from typing_extensions import override
 
 from pterodactyl import mcsrvstatus
 from pterodactyl.config import config, register_config
@@ -22,7 +23,7 @@ class Pterodactyl(commands.Cog):
 
     __author__ = ["[cswimr](https://www.coastalcommits.com/cswimr)"]
     __git__ = "https://www.coastalcommits.com/cswimr/SeaCogs"
-    __version__ = "2.0.5"
+    __version__ = "2.0.6"
     __documentation__ = "https://seacogs.coastalcommits.com/pterodactyl/"
 
     def __init__(self, bot: Red):
@@ -32,9 +33,10 @@ class Pterodactyl(commands.Cog):
         self.websocket: Optional[websockets.ClientConnection] = None
         self.retry_counter: int = 0
         register_config(config)
-        self.task = self.get_task()
+        self.task = self._get_task()
         self.update_topic.start()
 
+    @override
     def format_help_for_context(self, ctx: commands.Context) -> str:
         pre_processed = super().format_help_for_context(ctx) or ""
         n = "\n" if "\n\n" not in pre_processed else ""
@@ -46,50 +48,57 @@ class Pterodactyl(commands.Cog):
         ]
         return "\n".join(text)
 
+    @override
     async def cog_load(self) -> None:
         pterodactyl_keys = await self.bot.get_shared_api_tokens("pterodactyl")
         api_key = pterodactyl_keys.get("api_key")
         if api_key is None:
-            self.task.cancel()
+            self.maybe_cancel_task()
             logger.error("Pterodactyl API key not set. Please set it using `[p]set api`.")
             return
         base_url = await config.base_url()
         if base_url is None:
-            self.task.cancel()
+            self.maybe_cancel_task()
             logger.error("Pterodactyl base URL not set. Please set it using `[p]pterodactyl config url`.")
             return
         server_id = await config.server_id()
         if server_id is None:
-            self.task.cancel()
+            self.maybe_cancel_task()
             logger.error("Pterodactyl server ID not set. Please set it using `[p]pterodactyl config serverid`.")
             return
 
         self.client = PterodactylClient(base_url, api_key).client
 
+    @override
     async def cog_unload(self) -> None:
         self.update_topic.cancel()
-        self.task.cancel()
-        self.retry_counter = 0
+        self.maybe_cancel_task()
 
-    def get_task(self) -> asyncio.Task:
+    def maybe_cancel_task(self, reset_retry_counter: bool = True) -> None:
+        if self.task:
+            self.task.cancel()
+        if reset_retry_counter:
+            self.retry_counter = 0
+
+    def _get_task(self) -> asyncio.Task:
         from pterodactyl.websocket import establish_websocket_connection
 
         task = self.bot.loop.create_task(establish_websocket_connection(self), name="Pterodactyl Websocket Connection")
-        task.add_done_callback(self.error_callback)
+        task.add_done_callback(self._error_callback)
         return task
 
-    def error_callback(self, fut) -> None:  # NOTE Thanks flame442 and zephyrkul for helping me figure this out
+    def _error_callback(self, fut) -> None:  # NOTE Thanks flame442 and zephyrkul for helping me figure this out
         try:
             fut.result()
         except asyncio.CancelledError:
             logger.info("WebSocket task has been cancelled.")
         except Exception as e:  # pylint: disable=broad-exception-caught
             logger.error("WebSocket task has failed: %s", e, exc_info=e)
-            self.task.cancel()
+            self.maybe_cancel_task(reset_retry_counter=False)
             if self.retry_counter < 5:
                 self.retry_counter += 1
                 logger.info("Retrying in %s seconds...", 5 * self.retry_counter)
-                self.task = self.bot.loop.call_later(5 * self.retry_counter, self.get_task)
+                self.task = self.bot.loop.call_later(5 * self.retry_counter, self._get_task)
             else:
                 logger.info("Retry limit reached. Stopping task.")
 
@@ -100,9 +109,9 @@ class Pterodactyl(commands.Cog):
         console = self.bot.get_channel(await config.console_channel())
         chat = self.bot.get_channel(await config.chat_channel())
         if console:
-            await console.edit(topic=topic)
+            await console.edit(topic=topic)  # type: ignore
         if chat:
-            await chat.edit(topic=topic)
+            await chat.edit(topic=topic)  # type: ignore
 
     @commands.Cog.listener()
     async def on_message_without_command(self, message: discord.Message) -> None:
@@ -113,13 +122,7 @@ class Pterodactyl(commands.Cog):
                 return
             logger.debug("Received console command from %s: %s", message.author.id, message.content)
             await message.channel.send(f"Received console command from {message.author.id}: {message.content[:1900]}", allowed_mentions=discord.AllowedMentions.none())
-            try:
-                await self.websocket.send(json.dumps({"event": "send command", "args": [message.content]}))
-            except websockets.exceptions.ConnectionClosed as e:
-                logger.error("WebSocket connection closed: %s", e)
-                self.task.cancel()
-                self.retry_counter = 0
-                self.task = self.get_task()
+            await self._send(json.dumps({"event": "send command", "args": [message.content]}))
         if message.channel.id == await config.chat_channel() and message.author.bot is False:
             logger.debug("Received chat message from %s: %s", message.author.id, message.content)
             channel = self.bot.get_channel(await config.console_channel())
@@ -127,13 +130,22 @@ class Pterodactyl(commands.Cog):
                 await channel.send(f"Received chat message from {message.author.id}: {message.content[:1900]}", allowed_mentions=discord.AllowedMentions.none())
             msg = json.dumps({"event": "send command", "args": [await self.get_chat_command(message)]})
             logger.debug("Sending chat message to server:\n%s", msg)
+            await self._send(message=msg)
+
+    async def _send(self, message: Union[websockets.Data, Iterable[websockets.Data], AsyncIterable[websockets.Data]], text: bool = False):
+        """Send a message through the websocket connection. Restarts the websocket connection task if it is closed, and reinvokes itself."""
+        try:
+            await self.websocket.send(message=message, text=text)  # type: ignore - we want this to error if `self.websocket` is none
+        except websockets.exceptions.ConnectionClosed as e:
+            logger.error("WebSocket connection closed: %s", e)
+            self.maybe_cancel_task()
+            self.task = self._get_task()
             try:
-                await self.websocket.send(msg)
-            except websockets.exceptions.ConnectionClosed as e:
-                logger.error("WebSocket connection closed: %s", e)
-                self.task.cancel()
-                self.retry_counter = 0
-                self.task = self.get_task()
+                await asyncio.wait_for(fut=self.task, timeout=60)
+                await self._send(message=message, text=text)
+            except asyncio.TimeoutError:
+                logger.error("Timeout while waiting for websocket connection")
+                raise
 
     async def get_topic(self) -> str:
         topic: str = await config.topic()
@@ -193,7 +205,7 @@ class Pterodactyl(commands.Cog):
 
     async def get_player_list_embed(self, ctx: Union[commands.Context, discord.Interaction]) -> Optional[discord.Embed]:
         player_list = await self.get_player_list()
-        if player_list:
+        if player_list and isinstance(ctx.channel, discord.abc.Messageable):
             embed = discord.Embed(color=await self.bot.get_embed_color(ctx.channel), title="Players Online")
             embed.description = player_list[0]
             return embed
@@ -206,10 +218,12 @@ class Pterodactyl(commands.Cog):
         current_status = await config.current_status()
 
         if current_status == action_ing:
-            return await ctx.send(f"Server is already {action_ing}.", ephemeral=True)
+            await ctx.send(f"Server is already {action_ing}.", ephemeral=True)
+            return
 
         if current_status in ["starting", "stopping"] and action != "kill":
-            return await ctx.send("Another power action is already in progress.", ephemeral=True)
+            await ctx.send("Another power action is already in progress.", ephemeral=True)
+            return
 
         view = ConfirmView(ctx.author, disable_buttons=True)
 
@@ -220,13 +234,13 @@ class Pterodactyl(commands.Cog):
         if view.result is True:
             await message.edit(content=f"Sending websocket command to {action} server...", view=None)
 
-            await self.websocket.send(json.dumps({"event": "set state", "args": [action]}))
+            await self._websocket_send(json.dumps({"event": "set state", "args": [action]}))
 
             await message.edit(content=f"Server {action_ing}", view=None)
-            return None
+            return
 
         await message.edit(content="Cancelled.", view=None)
-        return None
+        return
 
     async def send_command(self, ctx: Union[discord.Interaction, commands.Context], command: str):
         channel = self.bot.get_channel(await config.console_channel())
@@ -234,23 +248,15 @@ class Pterodactyl(commands.Cog):
             ctx = await self.bot.get_context(ctx)
         if channel:
             await channel.send(f"Received console command from {ctx.author.id}: {command[:1900]}", allowed_mentions=discord.AllowedMentions.none())
-        try:
-            await self.websocket.send(json.dumps({"event": "send command", "args": [command]}))
-            await ctx.send(f"Command sent to server. {box(command, 'json')}")
-        except websockets.exceptions.ConnectionClosed as e:
-            logger.error("WebSocket connection closed: %s", e)
-            await ctx.send(error("WebSocket connection closed."))
-            self.task.cancel()
-            self.retry_counter = 0
-            self.task = self.get_task()
+        await self._websocket_send(json.dumps({"event": "send command", "args": [command]}))
+        await ctx.send(f"Command sent to server. {box(command, 'json')}")
 
     @commands.Cog.listener()
     async def on_red_api_tokens_update(self, service_name: str, api_tokens: Mapping[str, str]):  # pylint: disable=unused-argument
         if service_name == "pterodactyl":
             logger.info("Configuration value set: api_key\nRestarting task...")
-            self.task.cancel()
-            self.retry_counter = 0
-            self.task = self.get_task()
+            self.maybe_cancel_task(reset_retry_counter=True)
+            self.task = self._get_task()
 
     slash_pterodactyl = app_commands.Group(name="pterodactyl", description="Pterodactyl allows you to manage your Pterodactyl Panel from Discord.")
 
@@ -346,9 +352,8 @@ class Pterodactyl(commands.Cog):
         await config.base_url.set(base_url)
         await ctx.send(f"Base URL set to {base_url}")
         logger.info("Configuration value set: base_url = %s\nRestarting task...", base_url)
-        self.task.cancel()
-        self.retry_counter = 0
-        self.task = self.get_task()
+        self.maybe_cancel_task(reset_retry_counter=True)
+        self.task = self._get_task()
 
     @pterodactyl_config.command(name="serverid")
     async def pterodactyl_config_server_id(self, ctx: commands.Context, *, server_id: str) -> None:
@@ -356,9 +361,8 @@ class Pterodactyl(commands.Cog):
         await config.server_id.set(server_id)
         await ctx.send(f"Server ID set to {server_id}")
         logger.info("Configuration value set: server_id = %s\nRestarting task...", server_id)
-        self.task.cancel()
-        self.retry_counter = 0
-        self.task = self.get_task()
+        self.maybe_cancel_task(reset_retry_counter=True)
+        self.task = self._get_task()
 
     @pterodactyl_config.group(name="console")
     async def pterodactyl_config_console(self, ctx: commands.Context):

pterodactyl/websocket.py

@@ -2,7 +2,7 @@
 import json
 import re
 from pathlib import Path
-from typing import Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 import aiohttp
 import discord
@@ -56,7 +56,9 @@ async def establish_websocket_connection(coginstance: Pterodactyl) -> None:
                         content = mask_ip(content)
 
                     console_channel = coginstance.bot.get_channel(await config.console_channel())
+                    assert isinstance(console_channel, discord.abc.Messageable)
                     chat_channel = coginstance.bot.get_channel(await config.chat_channel())
+                    assert isinstance(chat_channel, discord.abc.Messageable)
                     if console_channel is not None:
                         if content.startswith("["):
                             pagified_content = pagify(content, delims=[" ", "\n"])
@@ -83,7 +85,7 @@ async def establish_websocket_connection(coginstance: Pterodactyl) -> None:
                                 embed, img = await generate_join_leave_embed(coginstance=coginstance, username=join_message, join=True)
                                 if img:
                                     with open(img, "rb") as file:
-                                        await chat_channel.send(embed=embed, file=file)
+                                        await chat_channel.send(embed=embed, file=discord.File(fp=file))
                                 else:
                                     await chat_channel.send(embed=embed)
                             else:
@@ -96,7 +98,7 @@ async def establish_websocket_connection(coginstance: Pterodactyl) -> None:
                                 embed, img = await generate_join_leave_embed(coginstance=coginstance, username=leave_message, join=False)
                                 if img:
                                     with open(img, "rb") as file:
-                                        await chat_channel.send(embed=embed, file=file)
+                                        await chat_channel.send(embed=embed, file=discord.File(fp=file))
                                 else:
                                     await chat_channel.send(embed=embed)
                             else:
@@ -106,7 +108,11 @@ async def establish_websocket_connection(coginstance: Pterodactyl) -> None:
                     if achievement_message:
                         if chat_channel is not None:
                             if coginstance.bot.embed_requested(chat_channel):
-                                await chat_channel.send(embed=await generate_achievement_embed(coginstance, achievement_message["username"], achievement_message["achievement"], achievement_message["challenge"]))
+                                embed, img = await generate_achievement_embed(coginstance, achievement_message["username"], achievement_message["achievement"], achievement_message["challenge"])
+                                if img:
+                                    await chat_channel.send(embed=embed, file=discord.File(fp=img))
+                                else:
+                                    await chat_channel.send(embed=embed)
                             else:
                                 await chat_channel.send(f"{achievement_message['username']} has {'completed the challenge' if achievement_message['challenge'] else 'made the advancement'} {achievement_message['achievement']}")
 
@@ -130,24 +136,27 @@ async def establish_websocket_connection(coginstance: Pterodactyl) -> None:
                                 await chat.send(await config.shutdown_msg())
 
 
-async def retrieve_websocket_credentials(coginstance: Pterodactyl) -> Optional[dict]:
+async def retrieve_websocket_credentials(coginstance: Pterodactyl) -> dict:
     pterodactyl_keys = await coginstance.bot.get_shared_api_tokens("pterodactyl")
     api_key = pterodactyl_keys.get("api_key")
     if api_key is None:
-        coginstance.task.cancel()
+        coginstance.maybe_cancel_task()
         raise ValueError("Pterodactyl API key not set. Please set it using `[p]set api`.")
     base_url = await config.base_url()
     if base_url is None:
-        coginstance.task.cancel()
+        coginstance.maybe_cancel_task()
         raise ValueError("Pterodactyl base URL not set. Please set it using `[p]pterodactyl config url`.")
     server_id = await config.server_id()
     if server_id is None:
-        coginstance.task.cancel()
+        coginstance.maybe_cancel_task()
         raise ValueError("Pterodactyl server ID not set. Please set it using `[p]pterodactyl config serverid`.")
 
     client = PterodactylClient(base_url, api_key).client
     coginstance.client = client
-    websocket_credentials = client.servers.get_websocket(server_id)
+    websocket_credentials: dict[str, Any] = client.servers.get_websocket(server_id).json()
+    if not websocket_credentials:
+        coginstance.maybe_cancel_task()
+        raise ValueError("Failed to retrieve websocket credentials. Please ensure the API details are correctly configured.")
     logger.debug(
         """Websocket connection details retrieved:
                         Socket: %s
@@ -165,44 +174,44 @@ def remove_ansi_escape_codes(text: str) -> str:
     return ansi_escape.sub("", text)
 
 
-async def check_if_server_message(text: str) -> Union[bool, str]:
+async def check_if_server_message(text: str) -> Optional[str]:
     regex = await config.server_regex()
     match: Optional[re.Match[str]] = re.match(regex, text)
     if match:
         logger.trace("Message is a server message")
         return match.group(1)
-    return False
+    return None
 
 
-async def check_if_chat_message(text: str) -> Union[bool, dict]:
+async def check_if_chat_message(text: str) -> Optional[dict]:
     regex = await config.chat_regex()
     match: Optional[re.Match[str]] = re.match(regex, text)
     if match:
         groups = {"username": match.group(1), "message": match.group(2)}
         logger.trace("Message is a chat message\n%s", json.dumps(groups))
         return groups
-    return False
+    return None
 
 
-async def check_if_join_message(text: str) -> Union[bool, str]:
+async def check_if_join_message(text: str) -> Optional[str]:
     regex = await config.join_regex()
     match: Optional[re.Match[str]] = re.match(regex, text)
     if match:
         logger.trace("Message is a join message")
         return match.group(1)
-    return False
+    return None
 
 
-async def check_if_leave_message(text: str) -> Union[bool, str]:
+async def check_if_leave_message(text: str) -> Optional[str]:
     regex = await config.leave_regex()
     match: Optional[re.Match[str]] = re.match(regex, text)
     if match:
         logger.trace("Message is a leave message")
         return match.group(1)
-    return False
+    return None
 
 
-async def check_if_achievement_message(text: str) -> Union[bool, dict]:
+async def check_if_achievement_message(text: str) -> Optional[dict]:
     regex = await config.achievement_regex()
     match: Optional[re.Match[str]] = re.match(regex, text)
     if match:
@@ -213,7 +222,7 @@ async def check_if_achievement_message(text: str) -> Union[bool, dict]:
             groups["challenge"] = False
         logger.trace("Message is an achievement message")
         return groups
-    return False
+    return None
 
 
 async def get_info(username: str) -> Optional[dict]: