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
This commit is contained in:
parent
ea0b7937f8
commit
2a5b924409
11 changed files with 184 additions and 139 deletions
|
@ -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):
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue