feat(repo): make all cogs pylance-typechecking compliant
Some checks failed
Actions / Lint Code (Ruff & Pylint) (push) Failing after 43s
Actions / Build Documentation (MkDocs) (push) Failing after 24s

at `basic` level, does not include Aurora as it's being rewritten in the `aurora/v3` branch
This commit is contained in:
cswimr 2025-02-01 16:57:45 +00:00
parent ea0b7937f8
commit 2a5b924409
Signed by: cswimr
GPG key ID: 0EC431A8DA8F8087
11 changed files with 184 additions and 139 deletions

View file

@ -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):