From d236f88b64575ea265477acd1a973166c887b2ab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Roman=20Kr=C4=8Dek?= <roman.krcek@tietoevry.com>
Date: Sun, 13 Oct 2024 19:47:06 +0200
Subject: [PATCH] Move to pydantic settings paradigm

---
 requirements.txt                    |  3 +-
 telegram_downloader_bot/main.py     | 25 ++++++--------
 telegram_downloader_bot/security.py |  6 ++--
 telegram_downloader_bot/settings.py | 53 +++++++++++++++++++++++++++++
 4 files changed, 68 insertions(+), 19 deletions(-)
 create mode 100644 telegram_downloader_bot/settings.py

diff --git a/requirements.txt b/requirements.txt
index fb40e1d..a2dd334 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,4 +5,5 @@ pyrogram==2.0.106
 tiktok_downloader==0.3.5
 uvloop==0.19.0
 tgcrypto==1.2.5
-sentry-sdk==2.15.0
\ No newline at end of file
+sentry-sdk==2.15.0
+pydantic-settings==2.5.2
\ No newline at end of file
diff --git a/telegram_downloader_bot/main.py b/telegram_downloader_bot/main.py
index fa53269..acc611f 100644
--- a/telegram_downloader_bot/main.py
+++ b/telegram_downloader_bot/main.py
@@ -1,4 +1,3 @@
-import os
 import uvloop
 
 from pyrogram import Client, filters
@@ -7,20 +6,18 @@ from pyrogram.types import Message
 from telegram_downloader_bot.logger import log
 from telegram_downloader_bot.telemetry import init_telemetry
 from telegram_downloader_bot import utils, security
-
-API_ID = os.getenv("API_ID")  # Your API ID from my.telegram.org
-API_HASH = os.getenv("API_HASH")  # Your API Hash from my.telegram.org
-BOT_TOKEN = os.getenv("BOT_TOKEN")  # Your bot token from BotFather
-STORAGE = os.getenv("STORAGE")  # Your bot token from BotFather
+from telegram_downloader_bot.settings import settings
 
 uvloop.install()
-init_telemetry()
+
+if settings.app_env == "PROD":
+    init_telemetry()
 
 app = Client("downloader_bot",
-             api_id=API_ID,
-             api_hash=API_HASH,
-             bot_token=BOT_TOKEN,
-             workers=1)
+             api_id=settings.api_id,
+             api_hash=settings.api_hash,
+             bot_token=settings.bot_token,
+             workers=settings.workers)
 
 
 @app.on_message(filters.command("start"))
@@ -59,7 +56,7 @@ async def message_handler(_, message: Message):
         msg = f"Downloading video {i+1}/{len(urls)}..."
         log.info(msg)
         await message.reply_text(msg)
-        utils.download_tt_video(STORAGE, url)
+        utils.download_tt_video(settings.storage, url)
 
     await message.reply_text("Done.")
 
@@ -70,9 +67,9 @@ async def media_handler(client, message: Message):
 
     await message.reply_text("Downloading media...")
 
-    utils.handle_media_message_contents(STORAGE, client, message)
+    utils.handle_media_message_contents(settings.storage, client, message)
 
 
 if __name__ == "__main__":
-    utils.make_fs(STORAGE)
+    utils.make_fs(settings.storage)
     app.run()
diff --git a/telegram_downloader_bot/security.py b/telegram_downloader_bot/security.py
index 06f5371..6c4ccc8 100644
--- a/telegram_downloader_bot/security.py
+++ b/telegram_downloader_bot/security.py
@@ -1,11 +1,9 @@
-import os
 from functools import wraps
 
 from telegram_downloader_bot.logger import log
+from telegram_downloader_bot.settings import settings
 
-# Comma separated list of Telegram IDs that this bot will respond to
-allowed_ids_raw = os.getenv("ALLOWED_IDS", "")
-allowed_ids = allowed_ids_raw.split(",")
+allowed_ids = settings.allowed_ids.split(",")
 allowed_ids = [int(x) for x in allowed_ids]
 
 
diff --git a/telegram_downloader_bot/settings.py b/telegram_downloader_bot/settings.py
new file mode 100644
index 0000000..8d3511d
--- /dev/null
+++ b/telegram_downloader_bot/settings.py
@@ -0,0 +1,53 @@
+import os
+from pydantic_settings import BaseSettings
+
+
+class Settings(BaseSettings):
+    """
+    Settings class that defines configuration variables for the application.
+
+    Attributes:
+    ----------
+    app_env : str
+        Specifies the environment in which the application is running.
+        Default is 'DEV'. Possible values could include 'DEV', 'PROD'
+
+    workers : int
+        Defines the number of workers to be used in the application.
+        Default is 1.
+
+    api_id : int
+        Represents the API ID from my.telegram.org
+
+    api_hash : str
+        The hash key corresponding to your API Hash from my.telegram.org
+
+    bot_token : str
+        The token from BotFather.
+
+    storage : os.path
+        Specifies the path where the application stores persistent data.
+
+    allowed_ids : str
+        A list or comma-separated string of IDs that are allowed access
+        to the bot or application.
+
+    Config:
+    -------
+    env_file : str
+        Specifies the environment file to load the environment variables from.
+        Default is ".env".
+    """
+    app_env: str = "DEV"
+    workers: int = 1
+    api_id: int
+    api_hash: str
+    bot_token: str
+    storage: os.path
+    allowed_ids: str
+
+    class Config:
+        env_file = ".env"
+
+
+settings = Settings()