update(launcher): improve launcher UI & shutdown handling

This commit is contained in:
Romain J 2020-08-28 23:05:04 +02:00
parent 9a0786af7c
commit 331599eb38
3 changed files with 160 additions and 79 deletions

View File

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

View File

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

View File

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