From 331599eb38eafaeeab0af8a599d36169029e4be3 Mon Sep 17 00:00:00 2001 From: Romain J Date: Fri, 28 Aug 2020 23:05:04 +0200 Subject: [PATCH] update(launcher): improve launcher UI & shutdown handling --- .idea/dictionaries/romain.xml | 3 + tuxbot/__main__.py | 68 +++++++------- tuxbot/core/bot.py | 168 +++++++++++++++++++++++++--------- 3 files changed, 160 insertions(+), 79 deletions(-) diff --git a/.idea/dictionaries/romain.xml b/.idea/dictionaries/romain.xml index e0976d7..2d9fb48 100644 --- a/.idea/dictionaries/romain.xml +++ b/.idea/dictionaries/romain.xml @@ -6,14 +6,17 @@ francais ipinfo iplocalise + jishaku levelname localiseip postgresql + releaselevel socketstats splt systemd tutux tuxbot + tuxbot's tuxvenv webhooks diff --git a/tuxbot/__main__.py b/tuxbot/__main__.py index 13d7a7c..a97747d 100644 --- a/tuxbot/__main__.py +++ b/tuxbot/__main__.py @@ -1,9 +1,7 @@ import argparse import asyncio -import getpass import json import logging -import platform import signal import sys import os @@ -12,7 +10,7 @@ from typing import NoReturn import discord import pip -from pip._vendor import distro +import tracemalloc from rich.columns import Columns from rich.console import Console from rich.panel import Panel @@ -30,6 +28,7 @@ log = logging.getLogger("tuxbot.main") console = Console() install(console=console) +tracemalloc.start() def list_instances() -> NoReturn: @@ -64,11 +63,11 @@ def list_instances() -> NoReturn: console.print(columns) console.print() - sys.exit(0) + sys.exit(os.EX_OK) def debug_info() -> NoReturn: - """Show debug infos relatives to the bot + """Show debug info relatives to the bot """ python_version = sys.version.replace("\n", "") @@ -76,12 +75,6 @@ def debug_info() -> NoReturn: tuxbot_version = __version__ dpy_version = discord.__version__ - os_info = distro.linux_distribution() - os_info = f"{os_info[0]} {os_info[1]}" - - runner = getpass.getuser() - - uname = os.popen('uname -a').read().strip().split() uptime = os.popen('uptime').read().strip().split() console.print( @@ -127,17 +120,20 @@ def debug_info() -> NoReturn: table.add_column( "Server Info", ) - table.add_row(f"[u]OS:[/u] {os_info}") - table.add_row(f"[u]Kernel:[/u] {uname[2]}") - table.add_row(f"[u]System arch:[/u] {platform.machine()}") - table.add_row(f"[u]User:[/u] {runner}") + table.add_row(f"[u]System:[/u] {os.uname().sysname}") + table.add_row(f"[u]System arch:[/u] {os.uname().machine}") + table.add_row(f"[u]Kernel:[/u] {os.uname().release}") + table.add_row(f"[u]User:[/u] {os.getlogin()}") table.add_row(f"[u]Uptime:[/u] {uptime[2]}") - table.add_row(f"[u]Load Average:[/u] {' '.join(uptime[-3:])}") + table.add_row( + f"[u]Load Average:[/u] {' '.join(map(str, os.getloadavg()))}" + ) columns.add_renderable(table) console.print(columns) console.print() - sys.exit(0) + + sys.exit(os.EX_OK) def parse_cli_flags(args: list) -> Namespace: @@ -159,8 +155,10 @@ def parse_cli_flags(args: list) -> Namespace: "--version", "-V", action="store_true", help="Show tuxbot's used version" ) - parser.add_argument("--debug", action="store_true", - help="Show debug information.") + parser.add_argument( + "--debug", action="store_true", + help="Show debug information." + ) parser.add_argument( "--list-instances", "-L", action="store_true", help="List all instance names" @@ -194,7 +192,6 @@ async def shutdown_handler(tux: Tux, signal_type, exit_code=None) -> NoReturn: """ if signal_type: log.info("%s received. Quitting...", signal_type) - sys.exit(ExitCodes.SHUTDOWN) elif exit_code is None: log.info("Shutting down from unhandled exception") tux.shutdown_code = ExitCodes.CRITICAL @@ -202,16 +199,7 @@ async def shutdown_handler(tux: Tux, signal_type, exit_code=None) -> NoReturn: if exit_code is not None: tux.shutdown_code = exit_code - try: - await tux.logout() - finally: - pending = [t for t in asyncio.all_tasks() if - t is not asyncio.current_task()] - - for task in pending: - task.cancel() - - await asyncio.gather(*pending, return_exceptions=True) + await tux.shutdown() async def run_bot(tux: Tux, cli_flags: Namespace) -> None: @@ -247,11 +235,17 @@ async def run_bot(tux: Tux, cli_flags: Namespace) -> None: try: await tux.load_packages() - await tux.start(token, bot=True) + console.print() + await tux.start(token=token, bot=True) except discord.LoginFailure: log.critical("This token appears to be valid.") - console.print_exception() + console.print() + console.print( + "[prompt.invalid]This token appears to be valid. [i]exiting...[/i]" + ) sys.exit(ExitCodes.CRITICAL) + except Exception as e: + raise e return None @@ -268,9 +262,10 @@ def main() -> NoReturn: elif cli_flags.debug: debug_info() elif cli_flags.version: - print("Tuxbot V3") + print(f"Tuxbot V{version_info.major}") print(f"Complete Version: {__version__}") - sys.exit(0) + + sys.exit(os.EX_OK) loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) @@ -292,10 +287,11 @@ def main() -> NoReturn: loop.run_until_complete(run_bot(tux, cli_flags)) except KeyboardInterrupt: console.print( - "[red]Please use quit instead of Ctrl+C to Shutdown!" + " [red]Please use quit instead of Ctrl+C to Shutdown!" ) log.warning("Please use quit instead of Ctrl+C to Shutdown!") log.error("Received KeyboardInterrupt") + console.print("[i]Trying to shutdown...") if tux is not None: loop.run_until_complete(shutdown_handler(tux, signal.SIGINT)) except SystemExit as exc: @@ -303,6 +299,7 @@ def main() -> NoReturn: if tux is not None: loop.run_until_complete(shutdown_handler(tux, None, exc.code)) except Exception as exc: + console.print_exception() log.exception("Unexpected exception (%s): ", type(exc), exc_info=exc) if tux is not None: loop.run_until_complete(shutdown_handler(tux, None, 1)) @@ -314,6 +311,7 @@ def main() -> NoReturn: loop.stop() loop.close() exit_code = ExitCodes.CRITICAL if tux is None else tux.shutdown_code + sys.exit(exit_code) diff --git a/tuxbot/core/bot.py b/tuxbot/core/bot.py index 813ae8a..995a93a 100644 --- a/tuxbot/core/bot.py +++ b/tuxbot/core/bot.py @@ -6,19 +6,24 @@ from typing import List, Union import discord from discord.ext import commands +from rich import box +from rich.columns import Columns +from rich.console import Console +from rich.panel import Panel +from rich.progress import Progress, TextColumn, BarColumn +from rich.table import Table from rich.traceback import install +from tuxbot import version_info from . import Config from .data_manager import logs_data_path -from .utils.functions.cli import bordered - from . import __version__, ExitCodes from .utils.functions.extra import ContextPlus - log = logging.getLogger("tuxbot") -install() +console = Console() +install(console=console) NAME = r""" _____ _ _ _ _ @@ -33,6 +38,13 @@ packages: List[str] = ["jishaku", "tuxbot.cogs.warnings", "tuxbot.cogs.admin"] class Tux(commands.AutoShardedBot): _loading: asyncio.Task + _progress = { + 'main': Progress( + TextColumn("[bold blue]{task.fields[task_name]}", justify="right"), + BarColumn() + ), + 'tasks': {} + } def __init__(self, *args, cli_flags=None, **kwargs): # by default, if the bot shutdown without any intervention, @@ -73,52 +85,86 @@ class Tux(commands.AutoShardedBot): async def load_packages(self): if packages: - print("Loading packages...") - for package in packages: - try: - self.load_extension(package) - except Exception as e: - print( - Fore.RED - + f"Failed to load package {package}" - + Style.RESET_ALL - + f" check " - f"{str((self.logs / 'tuxbot.log').resolve())} " - f"for more details" - ) + with Progress() as progress: + task = progress.add_task( + "Loading packages...", + total=len(packages) + ) - log.exception(f"Failed to load package {package}", exc_info=e) + for package in packages: + try: + self.load_extension(package) + progress.console.print(f"{package} loaded") + except Exception as e: + log.exception( + f"Failed to load package {package}", + exc_info=e + ) + progress.console.print( + f"[red]Failed to load package {package} " + f"[i](see " + f"{str((self.logs / 'tuxbot.log').resolve())} " + f"for more details)[/i]" + ) + + progress.advance(task) async def on_ready(self): self.uptime = datetime.datetime.now() - INFO = { - "title": "INFO", - "rows": [ - str(self.user), - f"Prefixes: {', '.join(self.config('core').get('prefixes'))}", - f"Language: {self.config('core').get('locale')}", - f"Tuxbot Version: {__version__}", - f"Discord.py Version: {discord.__version__}", - "Python Version: " + sys.version.replace("\n", ""), - f"Shards: {self.shard_count}", - f"Servers: {len(self.guilds)}", - f"Users: {len(self.users)}", - ], - } + self._progress.get("main").stop_task( + self._progress.get("tasks")["connecting"] + ) + self._progress.get("main").remove_task( + self._progress.get("tasks")["connecting"] + ) + console.clear() - COGS = {"title": "COGS", "rows": []} + console.print( + Panel(f"[bold blue]Tuxbot V{version_info.major}", style="blue"), + justify="center" + ) + console.print() + + columns = Columns(expand=True, padding=2, align="center") + + table = Table( + style="dim", border_style="not dim", + box=box.HEAVY_HEAD + ) + table.add_column( + "INFO", + ) + table.add_row(str(self.user)) + table.add_row(f"Prefixes: {', '.join(self.config('core').get('prefixes'))}") + table.add_row(f"Language: {self.config('core').get('locale')}") + table.add_row(f"Tuxbot Version: {__version__}") + table.add_row(f"Discord.py Version: {discord.__version__}") + table.add_row(f"Shards: {self.shard_count}") + table.add_row(f"Servers: {len(self.guilds)}") + table.add_row(f"Users: {len(self.users)}") + columns.add_renderable(table) + + table = Table( + style="dim", border_style="not dim", + box=box.HEAVY_HEAD + ) + table.add_column( + "COGS", + ) for extension in packages: - COGS["rows"].append( - f"[{'X' if extension in self.extensions else ' '}] {extension}" - ) + if extension in self.extensions: + status = f"[green]:heavy_check_mark: {extension} " + else: + status = f"[red]:cross_mark: {extension} " - print(Fore.LIGHTBLUE_EX + NAME) - print(Style.RESET_ALL) - print(bordered(INFO, COGS)) + table.add_row(status) + columns.add_renderable(table) - print(f"\n{'=' * 118}\n\n") + console.print(columns) + console.print() - async def is_owner(self, user: Union[discord.User, discord.Member]) -> bool: + async def is_owner(self, + user: Union[discord.User, discord.Member]) -> bool: """Determines if the user is a bot owner. Parameters @@ -154,9 +200,9 @@ class Tux(commands.AutoShardedBot): return if ( - message.guild.id in self.config.get_blacklist("guild") - or message.channel.id in self.config.get_blacklist("channel") - or message.author.id in self.config.get_blacklist("user") + message.guild.id in self.config.get_blacklist("guild") + or message.channel.id in self.config.get_blacklist("channel") + or message.author.id in self.config.get_blacklist("user") ): return @@ -170,11 +216,45 @@ class Tux(commands.AutoShardedBot): async def on_message(self, message: discord.Message): await self.process_commands(message) + async def start(self, token, bot): + """Connect to Discord and start all connections. + + Todo: add postgresql connect here + """ + with self._progress.get("main") as pg: + task_id = self._progress.get("tasks")["connecting"] = pg.add_task( + "connecting", + task_name="Connecting to Discord...", start=False + ) + pg.update(task_id) + await super().start(token, bot=bot) + async def logout(self): """Disconnect from Discord and closes all actives connections. Todo: add postgresql logout here """ + for task in self._progress.get("tasks").keys(): + self._progress.get("main").log("Shutting down", task) + + self._progress.get("main").stop_task( + self._progress.get("tasks")[task] + ) + self._progress.get("main").remove_task( + self._progress.get("tasks")["connecting"] + ) + self._progress.get("main").stop() + + pending = [ + t for t in asyncio.all_tasks() if + t is not asyncio.current_task() + ] + + for task in pending: + console.log("Canceling", task.get_name(), f"({task.get_coro()})") + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + await super().logout() async def shutdown(self, *, restart: bool = False):