Fix small bugs and added cleaner startup

Author: gccody

launcher.py

@@ -1,18 +1,37 @@
-import glob
 import atexit
+import glob
+import sys
+
+from time import sleep
+
+import numpy as np
+import requests
+
 from lib.bot import Bot
+from lib.config import Config
+from lib.progress import Progress
 
 COGS = [path.split("\\")[-1][:-3] if "\\" in path else path.split("/")[-1][:-3] for path in
         glob.glob('./lib/cogs/*.py')]
 VERSION = "0.0.1"
-bot: Bot = Bot(VERSION)
-print("Registering cogs...")
-for cog in COGS:
-    print(f" - {cog}")
-    bot.load_extension(f"lib.cogs.{cog}")
-print("All cogs ready, now registering shards...")
+config = Config()
+TOTAL_GUILDS = len(requests.get("https://discord.com/api/v9/users/@me/guilds",
+                                headers={"Authorization": f"Bot {config.token}"}).json())
+SHARDS_LIST: list[int] = list(range(round((TOTAL_GUILDS / 1000) + 1)))
+CLUSTERS = round((len(SHARDS_LIST) / 2) + 1)
+SHARDS_SPLIT: list[np.ndarray[int]] = np.array_split(SHARDS_LIST, CLUSTERS)
 
 
-if __name__ == '__main__':
+def start(cluster_id: int):
+    bot: Bot = Bot(list((SHARDS_SPLIT[cluster_id])), VERSION)
+    progress = Progress("Registering Cogs", len(COGS))
+    for index, cog in enumerate(COGS):
+        bot.load_extension(f"lib.cogs.{cog}")
+        sleep(.2)
+        progress.next()
     atexit.register(bot.db.commit)
     bot.run(bot.config.token)
+
+
+if __name__ == '__main__':
+    start(0)

lib/bot/__init__.py

@@ -21,25 +21,28 @@
 from discord import Interaction, Guild
 from lib.context import CustomContext
 from lib.urban import UrbanDictionary
+from lib.progress import Progress
 
 COMMAND_ERROR_REGEX = r"Command raised an exception: (.*?(?=: )): (.*)"
 IGNORE_EXCEPTIONS = (CommandNotFound, BadArgument, RuntimeError)
 
 
 class Bot(BaseBot):
-    def __init__(self, version: str):
+    def __init__(self, shards: list[int], version: str):
         self.config: Config = Config()
-
         self.db: DB = DB()
         self.db.build()
         self.cache: dict = dict()
         self.tasks = AsyncIOScheduler(timezone=str(tzlocal.get_localzone()))
+        self.logged_in = False
+        self.shards_ready = False
         self.ready = False
+        self.expected_shards = len(shards)
         self.urban: UrbanDictionary = UrbanDictionary()
         self.VERSION = version
         super().__init__(
                          owner_id="507214515641778187",
-                         shards=5000,
+                         shards=shards,
                          intents=discord.Intents.all(),
                          debug_guilds=["1064582321728143421"],
                          description="Misc Bot used for advanced moderation and guild customization")
@@ -48,14 +51,29 @@ def __init__(self, version: str):
     async def on_connect(self):
         self.tasks.start()
         await self.sync_commands()
-        print(f"Signed into {self.user.display_name}#{self.user.discriminator}\n")
-        self.ready = True
+        print(f"Signed into {self.user.display_name}#{self.user.discriminator}")
+        self.logged_in = True
+        if self.shards_ready and self.logged_in:
+            self.register_guilds()
 
     async def on_shard_ready(self, shard_id):
         await self.change_presence(activity=discord.Activity(type=discord.ActivityType.watching,
                                                              name=f"SHARD {shard_id}"),
                                    shard_id=shard_id)
-        print(f" - Shard {shard_id+1}/{self.shard_count} ready")
+        print(f"Shard Ready [{shard_id+1}/{self.shard_count}]")
+        if shard_id+1 == self.expected_shards:
+            self.shards_ready = True
+        if self.shards_ready and self.logged_in:
+            self.register_guilds()
+
+    def register_guilds(self):
+        progress = Progress('Registering guilds', len(self.guilds))
+        for guild in self.guilds:
+            self.db.execute("""INSERT OR IGNORE INTO guilds VALUES (?,?,?)""", guild.id, None, None)
+            progress.next()
+        self.db.commit()
+        self.ready = True
+        print('\nEverything is ready!')
 
     async def get_application_context(
         self, interaction: Interaction, cls=CustomContext
@@ -134,7 +152,8 @@ async def on_application_command_error(self, ctx: CustomContext, exc) -> None:
         else:
             raise exc
 
-    async def cpu_percent(self, interval = None, *args, **kwargs):
+    @staticmethod
+    async def cpu_percent(interval=None, *args, **kwargs):
         python_process = psutil.Process(os.getpid())
         if interval is not None and interval > 0.0:
             python_process.cpu_percent(*args, **kwargs)

lib/cogs/Info.py

@@ -1,6 +1,5 @@
 import os
 
-import discord
 import psutil
 from discord.commands import SlashCommandGroup
 from discord.embeds import Embed

lib/cogs/Youtube.py

@@ -1,12 +1,10 @@
 import re
 
 from discord.ext.commands import Cog
-from discord.commands import slash_command, SlashCommandGroup, option
+from discord.commands import SlashCommandGroup, option
 from discord.embeds import Embed
 import aiohttp
 
-from bs4 import BeautifulSoup
-
 from lib.bot import Bot
 from lib.context import CustomContext
 

lib/db/__init__.py

@@ -1,7 +1,9 @@
 import re
 import sys
+import time
 from sqlite3 import Connection, Cursor, connect, Error
 
+from lib.progress import Progress
 
 class DB:
     DB_PATH = "./data/database.db"
@@ -96,15 +98,13 @@ def update_table(self, table: str, data: str):
         self.cur.execute(f"ALTER TABLE {table}_new RENAME TO {table};")
 
     def build(self) -> None:
-        print('Building Database...')
+        progress = Progress("Building Database", len(self.table_info))
         for table in self.table_info:
             table_data = table['table']
             for table_name, columns in table_data.items():
-                print(f' - {table_name}')
                 create_table_query = f"CREATE TABLE IF NOT EXISTS {table_name}("
                 constraints = table.get('constraints')
                 for column_name, data_type in columns.items():
-                    print(f'    - {column_name}')
                     if constraints:
                         not_null = constraints.get('not_null')
                         if not_null:
@@ -136,6 +136,8 @@ def build(self) -> None:
                 self.cur.execute(create_table_query)
                 data = re.search(r"\((?<=\()(?:.)+", create_table_query).group()
                 self.update_table(table_name, data)
+                progress.next()
+                time.sleep(.2)
         self.commit()
 
     def commit(self) -> None:

lib/progress/__init__.py

@@ -0,0 +1,17 @@
+import sys
+
+
+class Progress:
+    def __init__(self, msg: str, total: int):
+        self.msg = msg
+        self.iteration = 0
+        self.total = total
+        sys.stdout.write(f"\r{self.msg} [{self.iteration}/{self.total}]")
+        sys.stdout.flush()
+
+    def next(self):
+        self.iteration += 1
+        sys.stdout.write(f"\r{self.msg} [{self.iteration}/{self.total}]")
+        sys.stdout.flush()
+        if self.iteration == self.total:
+            print("")