Merge pull request #3 from slashtechno/restructure-project

Minimize code in in `__main__.py` by utilizing files under `utils/`

llmail/utils/__init__.py

@@ -0,0 +1,5 @@
+from llmail.utils.logging import set_primary_logger, logger
+from llmail.utils.cli_args import args, bot_email
+
+# https://stackoverflow.com/a/31079085
+__all__ = ["set_primary_logger", "logger", "args", bot_email]

llmail/utils/cli_args.py

@@ -1,10 +1,9 @@
 import argparse
 import os
-import dotenv
 import sys
 from pathlib import Path
 
-argparser = None
+import dotenv
 
 """
 For reference, Gmail's IMAP settings are:
@@ -16,7 +15,7 @@
 
 
 def set_argparse():
-    global argparser
+    global args
 
     if Path(".env").is_file():
         dotenv.load_dotenv()
@@ -37,9 +36,7 @@ def set_argparse():
         title="Subcommands",
     )
     # Subcommand: list-folders
-    _ = subparsers.add_parser(
-        "list-folders", help="List all folders in the IMAP account and exit"
-    )
+    _ = subparsers.add_parser("list-folders", help="List all folders in the IMAP account and exit")
     # General arguments
     argparser.add_argument(
         "--log-level",
@@ -66,9 +63,7 @@ def set_argparse():
         "-w",
         help="Interval in seconds to check for new emails. If not set, will only check once.",
         type=int,
-        default=(
-            int(os.getenv("WATCH_INTERVAL")) if os.getenv("WATCH_INTERVAL") else None
-        ),
+        default=(int(os.getenv("WATCH_INTERVAL")) if os.getenv("WATCH_INTERVAL") else None),
     )
     # OpenAI-compatible API arguments
     ai_api = argparser.add_argument_group("OpenAI-compatible API")
@@ -108,22 +103,16 @@ def set_argparse():
         "--subject-key",
         "-s",
         help="Emails with this subject will be replied to",
-        default=(
-            os.getenv("SUBJECT_KEY") if os.getenv("SUBJECT_KEY") else "llmail autoreply"
-        ),
+        default=(os.getenv("SUBJECT_KEY") if os.getenv("SUBJECT_KEY") else "llmail autoreply"),
     )
     email.add_argument(
         "--alias",
         help="Name to use in the 'From' in addition to the email address",
         default=os.getenv("ALIAS") if os.getenv("ALIAS") else "LLMail",
     )
     imap = email.add_argument_group("IMAP")
-    imap.add_argument(
-        "--imap-host", help="IMAP server hostname", default=os.getenv("IMAP_HOST")
-    )
-    imap.add_argument(
-        "--imap-port", help="IMAP server port", default=os.getenv("IMAP_PORT")
-    )
+    imap.add_argument("--imap-host", help="IMAP server hostname", default=os.getenv("IMAP_HOST"))
+    imap.add_argument("--imap-port", help="IMAP server port", default=os.getenv("IMAP_PORT"))
     imap.add_argument(
         "--imap-username",
         help="IMAP server username",
@@ -135,12 +124,8 @@ def set_argparse():
         default=os.getenv("IMAP_PASSWORD"),
     )
     smtp = email.add_argument_group("SMTP")
-    smtp.add_argument(
-        "--smtp-host", help="SMTP server hostname", default=os.getenv("SMTP_HOST")
-    )
-    smtp.add_argument(
-        "--smtp-port", help="SMTP server port", default=os.getenv("SMTP_PORT")
-    )
+    smtp.add_argument("--smtp-host", help="SMTP server hostname", default=os.getenv("SMTP_HOST"))
+    smtp.add_argument("--smtp-port", help="SMTP server port", default=os.getenv("SMTP_PORT"))
     smtp.add_argument(
         "--smtp-username",
         help="SMTP server username",
@@ -154,9 +139,7 @@ def set_argparse():
     smtp.add_argument(
         "--message-id-domain",
         help="Domain to use for Message-ID header",
-        default=(
-            os.getenv("MESSAGE_ID_DOMAIN") if os.getenv("MESSAGE_ID_DOMAIN") else None
-        ),
+        default=(os.getenv("MESSAGE_ID_DOMAIN") if os.getenv("MESSAGE_ID_DOMAIN") else None),
     )
 
     check_required_args(
@@ -174,6 +157,10 @@ def set_argparse():
         ],
         argparser,
     )
+    args = argparser.parse_args()
+    # Setting bot_email instead of using imap_username directly in case support is needed for imap_username and bot_email being different
+    global bot_email
+    bot_email = args.imap_username
 
 
 def check_required_args(required_args: list[str], argparser: argparse.ArgumentParser):

llmail/utils/logging.py

@@ -0,0 +1,27 @@
+import re
+from sys import stderr
+
+from loguru import logger
+
+from llmail.utils.cli_args import args
+
+logging_file = stderr
+
+
+def redact_email_sink(message: str):
+    """Custom sink function that redacts email addresses before logging."""
+    email_pattern = r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b"
+    redacted_message = re.sub(email_pattern, "[redacted]", message)
+    print(redacted_message, file=logging_file)
+
+
+def set_primary_logger(log_level, redact_email_addresses):
+    """Set up the primary logger with the specified log level. Output to stderr and use the format specified."""
+    logger.remove()
+    # ^10 is a formatting directive to center with a padding of 10
+    logger_format = "<green>{time:YYYY-MM-DD HH:mm:ss}</green> |<level>{level: ^10}</level>| <level>{message}</level>"
+    sink = redact_email_sink if redact_email_addresses else stderr
+    logger.add(sink=sink, format=logger_format, colorize=True, level=log_level)
+
+
+set_primary_logger(args.log_level, args.redact_email_addresses)