update(launcher): improve launcher UI & shutdown handling
This commit is contained in:
parent
9a0786af7c
commit
331599eb38
3 changed files with 160 additions and 79 deletions
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -296,6 +291,7 @@ def main() -> NoReturn:
|
|||
)
|
||||
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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
task = progress.add_task(
|
||||
"Loading packages...",
|
||||
total=len(packages)
|
||||
)
|
||||
|
||||
for package in packages:
|
||||
try:
|
||||
self.load_extension(package)
|
||||
progress.console.print(f"{package} loaded")
|
||||
except Exception as e:
|
||||
print(
|
||||
Fore.RED
|
||||
+ f"Failed to load package {package}"
|
||||
+ Style.RESET_ALL
|
||||
+ f" check "
|
||||
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"
|
||||
f"for more details)[/i]"
|
||||
)
|
||||
|
||||
log.exception(f"Failed to load package {package}", exc_info=e)
|
||||
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)}",
|
||||
],
|
||||
}
|
||||
|
||||
COGS = {"title": "COGS", "rows": []}
|
||||
for extension in packages:
|
||||
COGS["rows"].append(
|
||||
f"[{'X' if extension in self.extensions else ' '}] {extension}"
|
||||
self._progress.get("main").stop_task(
|
||||
self._progress.get("tasks")["connecting"]
|
||||
)
|
||||
self._progress.get("main").remove_task(
|
||||
self._progress.get("tasks")["connecting"]
|
||||
)
|
||||
console.clear()
|
||||
|
||||
print(Fore.LIGHTBLUE_EX + NAME)
|
||||
print(Style.RESET_ALL)
|
||||
print(bordered(INFO, COGS))
|
||||
console.print(
|
||||
Panel(f"[bold blue]Tuxbot V{version_info.major}", style="blue"),
|
||||
justify="center"
|
||||
)
|
||||
console.print()
|
||||
|
||||
print(f"\n{'=' * 118}\n\n")
|
||||
columns = Columns(expand=True, padding=2, align="center")
|
||||
|
||||
async def is_owner(self, user: Union[discord.User, discord.Member]) -> bool:
|
||||
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:
|
||||
if extension in self.extensions:
|
||||
status = f"[green]:heavy_check_mark: {extension} "
|
||||
else:
|
||||
status = f"[red]:cross_mark: {extension} "
|
||||
|
||||
table.add_row(status)
|
||||
columns.add_renderable(table)
|
||||
|
||||
console.print(columns)
|
||||
console.print()
|
||||
|
||||
async def is_owner(self,
|
||||
user: Union[discord.User, discord.Member]) -> bool:
|
||||
"""Determines if the user is a bot owner.
|
||||
|
||||
Parameters
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Reference in a new issue