Refactor main function into ChaosBot class for improved structure and initialization of components

botbotbot/__init__.py

@@ -1,5 +1,6 @@
 import logging
 import tomllib
+from typing import Any
 
 import discord
 
@@ -11,31 +12,55 @@ from botbotbot.tts import CambAI
 from botbotbot.voice import VoiceBot
 from botbotbot.wordlist import Wordlist
 
+logger = logging.getLogger(__name__)
 
-def main() -> None:
-    description = """Discord Chaos Bot"""
 
-    logger = logging.getLogger(__name__)
-    logging.basicConfig(level=logging.INFO)
+class ChaosBot:
+    def __init__(self, config: dict[str, Any]) -> None:
+        self.config = config
 
-    with open("config.toml", "rb") as config_file:
-        config = tomllib.load(config_file)
+        intents = discord.Intents.all()
+        intents.presences = False
 
-    guild_ids = config.get("guild_ids")
-    if not (isinstance(guild_ids, list) and all(isinstance(i, int) for i in guild_ids)):
-        logger.error("Guild IDs must be a list of integers.")
-        guild_ids = []
+        self.bot = discord.Bot(description="Discord Chaos Bot", intents=intents)
 
+        self.init_aibot()
+        self.init_cambai()
+        self.init_guild_ids()
+
+        self.shuffler = Shuffler()
+        self.word_list = Wordlist(self.bot, self.guild_ids)
+        self.text_bot = TextBot(
+            self.bot,
+            self.word_list,
+            aibot=self.ai_bot,
+            shuffler=self.shuffler,
+            guild_ids=self.guild_ids,
+        )
+        self.voice_bot = VoiceBot(
+            self.bot,
+            self.cambai,
+            aibot=self.ai_bot,
+            shuffler=self.shuffler,
+            guild_ids=self.guild_ids,
+        )
+        self.nick_bot = NickBot(
+            self.bot, shuffler=self.shuffler, guild_ids=self.guild_ids
+        )
+
+        self.init_events()
+
+    def init_aibot(self) -> None:
+        self.ai_bot: AIBot | None = None
+
+        if isinstance(key := self.config.get("mistral_api_key"), str):
             system_prompt = """Tu es une intelligence artificelle qui répond en français.
             Tu dois faire un commentaire pertinent en lien avec ce qui te sera dit.
             Ta réponse doit être très courte.
             Ta réponse doit être une seule phrase.
             TA RÉPONSE DOIT ÊTRE EN FRANÇAIS !!!"""
 
-    aibot: AIBot | None = None
-
-    if isinstance(key := config.get("mistral_api_key"), str):
-        aibot = AIBot(
+            self.ai_bot = AIBot(
                 key,
                 model="mistral-large-latest",
                 system_message=system_prompt,
@@ -43,34 +68,42 @@ def main() -> None:
         else:
             logger.warning("No AI Bot.")
 
-    cambai: CambAI | None = None
-    if isinstance(key := config.get("cambai_api_key"), str):
-        cambai = CambAI(key)
+    def init_cambai(self) -> None:
+        self.cambai: CambAI | None = None
+        if isinstance(key := self.config.get("cambai_api_key"), str):
+            self.cambai = CambAI(key)
         else:
             logger.warning("No CambAI.")
 
-    intents = discord.Intents.default()
-    intents.members = True
-    intents.message_content = True
-    intents.reactions = True
-    intents.voice_states = True
+    def init_guild_ids(self) -> None:
+        guild_ids = self.config.get("guild_ids")
+        if not (
+            isinstance(guild_ids, list) and all(isinstance(i, int) for i in guild_ids)
+        ):
+            logger.error("Guild IDs must be a list of integers.")
+            guild_ids = []
+        self.guild_ids = guild_ids
 
-    bot = discord.Bot(description=description, intents=intents)
-    shf = Shuffler()
-    wl = Wordlist(bot, guild_ids)
-    wl.init_events()
+    def init_events(self) -> None:
+        self.bot.add_listener(self.on_ready, "on_ready")
 
-    text_bot = TextBot(bot, wl, aibot=aibot, shuffler=shf, guild_ids=guild_ids)
-    text_bot.init_events()
+        self.word_list.init_events()
+        self.text_bot.init_events()
+        self.voice_bot.init_events()
+        self.nick_bot.init_events()
 
-    voice_bot = VoiceBot(bot, cambai, aibot=aibot, shuffler=shf, guild_ids=guild_ids)
-    voice_bot.init_events()
+    async def on_ready(self) -> None:
+        logger.info(f"We have logged in as {self.bot.user}")
 
-    nick_bot = NickBot(bot, shuffler=shf, guild_ids=guild_ids)
-    nick_bot.init_events()
+    def run(self) -> None:
+        self.bot.run(self.config.get("token"))
 
-    @bot.listen("on_ready")
-    async def on_ready() -> None:
-        logger.info(f"We have logged in as {bot.user}")
 
-    bot.run(config.get("token"))
+def main() -> None:
+    logging.basicConfig(level=logging.INFO)
+
+    with open("config.toml", "rb") as config_file:
+        config = tomllib.load(config_file)
+
+    chaos_bot = ChaosBot(config)
+    chaos_bot.run()