from genericpath import isdir import logging import os import subprocess from typing import List, Optional import re from tempfile import NamedTemporaryFile from telegram import Update from telegram.ext import ( Updater, CommandHandler, DictPersistence, CallbackContext, defaults, ) from telegram.parsemode import ParseMode from parser import SafeArgumentParser, safe_str class KollagenBot: def __init__( self, tg_token: str, kollagen_path: str, base_dir: Optional[str] ) -> None: self.logger = logging.getLogger("kollagen") self.kollagen_path = kollagen_path self.base_dir = base_dir self._init_parser() self.updater = Updater( tg_token, persistence=DictPersistence(user_data_json="{}") ) dispatcher = self.updater.dispatcher dispatcher.add_handler(CommandHandler("start", self.tg_start)) dispatcher.add_handler(CommandHandler("help", self.tg_help)) dispatcher.add_handler(CommandHandler("generate", self.tg_generate)) dispatcher.add_handler(CommandHandler("g", self.tg_generate)) dispatcher.add_handler(CommandHandler("regenerate", self.tg_regenerate)) dispatcher.add_handler(CommandHandler("r", self.tg_regenerate)) dispatcher.add_error_handler(self.tg_error) def _init_parser(self): parser = SafeArgumentParser(prog="/generate", add_help=False) parser.add_argument( "directories", metavar="dir", type=safe_str, nargs="*", default=[self.base_dir] if self.base_dir else [], help="Directories to process. By default, the entire directory is processed.", ) parser.add_argument( "-m", dest="mode", metavar="mode", choices=self._get_modes(), nargs="?", const=True, help=f"Which collage mode to use. By default, one is chosen at random. When no value is specified, all modes are listed.", ) parser.add_argument( "-n", dest="num_images", metavar="N", type=int, help=f"How many images to have in a collage. Random (2 None: self.logger.error( msg="Exception while handling an update:", exc_info=context.error ) if isinstance(update, Update): if isinstance(context.error, subprocess.CalledProcessError): update.message.reply_text( f"Something is fucked!\n{context.error.stderr.decode('utf-8')}" ) else: update.message.reply_text(f"Something is fucked!\n{context.error}") def start_idle(self): self.updater.start_polling() self.updater.idle() def main() -> None: logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=getattr(logging, os.getenv("LOG_LEVEL", "info").upper()), ) tg_token = os.getenv("TG_TOKEN") if not tg_token: logging.error("TG_TOKEN is required.") exit(1) if env_kollagen_path := os.getenv("KOLLAGEN_PATH"): if os.path.exists(env_kollagen_path): kollagen_path = env_kollagen_path else: logging.error(f"kollagen not found! {env_kollagen_path} does not exist.") exit(1) else: which = subprocess.run(["which", "kollagen"], capture_output=True) try: which.check_returncode() kollagen_path = which.stdout.decode("utf-8").strip() except subprocess.CalledProcessError: logging.error( "kollagen not found! KOLLAGEN_PATH not specified and `kollagen` isn't in $PATH." ) exit(1) bot = KollagenBot(tg_token, kollagen_path, os.getenv("BASE_DIR")) bot.start_idle() if __name__ == "__main__": main()