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