update(launcher): improve launcher UI & shutdown handling

.idea/dictionaries/romain.xml

@@ -6,14 +6,17 @@
       <w>francais</w>
       <w>ipinfo</w>
       <w>iplocalise</w>
+      <w>jishaku</w>
       <w>levelname</w>
       <w>localiseip</w>
       <w>postgresql</w>
+      <w>releaselevel</w>
       <w>socketstats</w>
       <w>splt</w>
       <w>systemd</w>
       <w>tutux</w>
       <w>tuxbot</w>
+      <w>tuxbot's</w>
       <w>tuxvenv</w>
       <w>webhooks</w>
     </words>

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]System:[/u] {os.uname().sysname}")
-    table.add_row(f"[u]Kernel:[/u] {uname[2]}")
+    table.add_row(f"[u]System arch:[/u] {os.uname().machine}")
-    table.add_row(f"[u]System arch:[/u] {platform.machine()}")
+    table.add_row(f"[u]Kernel:[/u] {os.uname().release}")
-    table.add_row(f"[u]User:[/u] {runner}")
+    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",
+    parser.add_argument(
-                        help="Show debug information.")
+        "--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.shutdown()
-        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)
 
 
 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 <prefix>quit instead of Ctrl+C to Shutdown!"
+            "  [red]Please use <prefix>quit instead of Ctrl+C to Shutdown!"
         )
         log.warning("Please use <prefix>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)
 
 

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...")
+            with Progress() as progress:
-            for package in packages:
+                task = progress.add_task(
-                try:
+                    "Loading packages...",
-                    self.load_extension(package)
+                    total=len(packages)
-                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"
-                    )
 
-                    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 = {
+        self._progress.get("main").stop_task(
-            "title": "INFO",
+            self._progress.get("tasks")["connecting"]
-            "rows": [
+        )
-                str(self.user),
+        self._progress.get("main").remove_task(
-                f"Prefixes: {', '.join(self.config('core').get('prefixes'))}",
+            self._progress.get("tasks")["connecting"]
-                f"Language: {self.config('core').get('locale')}",
+        )
-                f"Tuxbot Version: {__version__}",
+        console.clear()
-                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)}",
-            ],
-        }
 
-        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(
+            if extension in self.extensions:
-                f"[{'X' if extension in self.extensions else ' '}] {extension}"
+                status = f"[green]:heavy_check_mark: {extension} "
-            )
+            else:
+                status = f"[red]:cross_mark: {extension} "
 
-        print(Fore.LIGHTBLUE_EX + NAME)
+            table.add_row(status)
-        print(Style.RESET_ALL)
+        columns.add_renderable(table)
-        print(bordered(INFO, COGS))
 
-        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")
+                message.guild.id in self.config.get_blacklist("guild")
-            or message.channel.id in self.config.get_blacklist("channel")
+                or message.channel.id in self.config.get_blacklist("channel")
-            or message.author.id in self.config.get_blacklist("user")
+                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):