Merge pull request #17 from YeetCode-devs/staging/hakimi

iSort, coloured logging, add_license.py improvement, module system, pyrogram

Author: prathamdby

modules/start.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: GPL-3.0-only
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Copyright (c) 2024, YeetCode Developers <[email protected]>
+
+from pyrogram import filters
+from pyrogram.client import Client
+from pyrogram.handlers import MessageHandler
+from pyrogram.types import Message
+
+from src.Module import ModuleBase
+
+
+class Module(ModuleBase):
+    def on_load(self, app: Client):
+        app.add_handler(MessageHandler(start, filters.command("start")))
+
+    def on_shutdown(self, app: Client):
+        pass
+
+
+async def start(app: Client, message: Message):
+    await message.reply("Hello!")

poetry.lock

@@ -10,8 +10,9 @@ package-mode = false
 [tool.poetry.dependencies]
 python = "^3.10"
 g4f = {extras = ["all"], version = "^0.2.5.4"}
-telethon = "^1.34.0"
 python-dotenv = "^1.0.1"
+pyrogram = {url = "https://github.com/KurimuzonAkuma/pyrogram/archive/refs/heads/dev.zip"}
+tgcrypto = "^1.2.5"
 
 [tool.poetry.group.dev.dependencies]
 black = "^24.3.0"

pyproject.toml

@@ -41,7 +41,7 @@
 
 
 def add_license_header(file_path) -> StatusIsOk:
-    with open(file_path, "r+") as f:
+    with open(file_path, "r+", encoding="utf-8") as f:
         content = f.read()
         if license_header in content:
             return False

scripts/add_license.py

@@ -17,11 +17,13 @@
 from os import getenv
 
 from dotenv import load_dotenv
-from telethon import TelegramClient, events
+from pyrogram.client import Client
+
+from .Module import load_modules
 
 
 def main() -> None:
-    load_dotenv()
+    load_dotenv(override=True)
 
     api_id = getenv("API_ID", 0)
     api_hash = getenv("API_HASH", "")
@@ -30,10 +32,7 @@ def main() -> None:
     if not all([api_id, api_hash, bot_token]):
         raise ValueError("Could not get all required credentials from env!")
 
-    app = TelegramClient("app", int(api_id), api_hash).start(bot_token=bot_token)
-
-    @app.on(events.NewMessage(incoming=True, pattern="/start"))
-    async def start(event):
-        await event.reply("Hello!")
+    app = Client("app", int(api_id), api_hash, bot_token=bot_token)
 
-    app.run_until_disconnected()
+    loaded_modules = load_modules(app)
+    app.run()

src/Bot.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: GPL-3.0-only
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Copyright (c) 2024, YeetCode Developers <[email protected]>
+
+import logging
+import os
+
+GLOBAL_DEBUG: bool = False
+if os.getenv("TGBOT_DEBUG") is not None:
+    GLOBAL_DEBUG = True
+
+log_additional_args: dict = {"filename": "bot.log", "level": logging.INFO}
+if GLOBAL_DEBUG:
+    log_additional_args.clear()
+    log_additional_args.update({"level": logging.DEBUG})
+
+
+#
+# Adapted from https://stackoverflow.com/a/56944256
+#
+class ColouredFormatter(logging.Formatter):
+
+    grey = "\x1b[38;20m"
+    yellow = "\x1b[33;20m"
+    red = "\x1b[31;20m"
+    green = "\x1b[0;32m"
+    blue = "\x1b[0;34m"
+    bold_red = "\x1b[31;1m"
+    reset = "\x1b[0m"
+    format_str = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+
+    FORMATS = {
+        logging.DEBUG: blue + format_str + reset,
+        logging.INFO: green + format_str + reset,
+        logging.WARNING: yellow + format_str + reset,
+        logging.ERROR: red + format_str + reset,
+        logging.CRITICAL: bold_red + format_str + reset,
+    }
+
+    def format(self, record: logging.LogRecord):
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt)
+        return formatter.format(record)
+
+
+logging.basicConfig(format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", **log_additional_args)
+
+
+for handler in logging.root.handlers:
+    if issubclass(logging.StreamHandler, type(handler)):
+        logging.root.removeHandler(handler)
+
+_sh = logging.StreamHandler()
+_sh.setFormatter(ColouredFormatter())
+logging.root.addHandler(_sh)
+logging.getLogger(__name__).info("Coloured log output initialized")

src/Logging.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: GPL-3.0-only
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Copyright (c) 2024, YeetCode Developers <[email protected]>
+
+import atexit
+import logging
+from abc import ABC, abstractmethod
+from importlib import import_module
+from pathlib import Path
+
+from pyrogram.client import Client
+
+log: logging.Logger = logging.getLogger(__name__)
+
+
+class ModuleBase(ABC):
+    @abstractmethod
+    def on_load(self, app: Client):
+        pass
+
+    @abstractmethod
+    def on_shutdown(self, app: Client):
+        pass
+
+
+def load_modules(app: Client) -> list[object]:
+    loaded_modules: list[object] = []
+
+    log.info("Searching for modules")
+    modules: list[Path] = list(Path("modules").rglob("*.py"))
+    log.info(f"Found {len(modules)} modules")
+
+    for module in modules:
+        log.info(f"Loading module '{module}'")
+
+        mdl = import_module(f"modules.{module.name.removesuffix('.py')}")
+
+        if not hasattr(mdl, "Module"):
+            log.error(f"Module '{module}' does not have a Module class, cannot load")
+            continue
+
+        if not issubclass(mdl.Module, ModuleBase):
+            log.warning(f"Module '{module}' does not inherit from ModuleBase class")
+
+        mdl.Module().on_load(app)
+        atexit.register(mdl.Module().on_shutdown, app)
+
+        loaded_modules.append(mdl)
+
+    return loaded_modules

src/Module.py

@@ -14,6 +14,7 @@
 #
 # Copyright (c) 2024, YeetCode Developers <[email protected]>
 
+from . import Logging
 from .Bot import main
 
 if __name__ == "__main__":