thm
/
kollagen
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

actual cmdline parsing in tg

master
Tomáš Mládek 2021-09-19 13:19:00 +02:00
parent d443fc6d7e
commit 8915a36b3d
2 changed files with 116 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

119
tgbot/kollagen-bot/main.py
View File

@ -2,18 +2,22 @@ from genericpath import isdir
import logging
import os
import subprocess
from typing import Optional
from typing import List, Optional
import re
from tempfile import NamedTemporaryFile
from telegram import Update, ForceReply
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:
@ -24,26 +28,78 @@ class KollagenBot:
self.kollagen_path = kollagen_path
self.base_dir = base_dir
self.updater = Updater(tg_token, persistence=DictPersistence(user_data_json="{}"))
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("list_modes", self.tg_list_modes))
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<n<5) by default.",
)
parser.add_argument(
"-w",
dest="width",
type=int,
default=640,
help=f"Width of resulting output (in px). 640px by default.",
)
parser.add_argument(
"-h",
dest="height",
type=int,
default=640,
help=f"Height of resulting output (in px). 640px by default.",
)
self.parser = parser
def _get_modes(self):
modes = subprocess.run(
[self.kollagen_path, "-m"], check=True, capture_output=True
)
return modes.stdout.decode("utf-8").strip().split(", ")
def tg_start(self, update: Update, context: CallbackContext):
update.message.reply_text("Hi! Check out https://gitlab.com/tmladek/kollagen")
self._process("", update)
self._process([], update)
def tg_generate(self, update: Update, context: CallbackContext):
cmd_line = update.message.text.split(" ")[1:]
self._process(cmd_line, update)
if context.user_data is not None:
success = self._process(cmd_line, update)
if success and context.user_data is not None:
context.user_data["last_cmd_line"] = cmd_line
def tg_regenerate(self, update: Update, context: CallbackContext):
@ -52,31 +108,56 @@ class KollagenBot:
else:
update.message.reply_text("No previous command to regenerate!")
def _process(self, cmd_line: str, update: Update):
def _process(self, cmd_line: List[str], update: Update):
self.logger.info(
f"Generating from {update.effective_user}, with cmd_line: `{cmd_line}`"
)
args = self.parser.parse_args(cmd_line)
if args.mode is True:
update.message.reply_text(
f"Available modes: {', '.join(self._get_modes())}"
)
return False
directories = [
os.path.join(self.base_dir or "./", re.sub(r"[^a-zA-Z0-9]", "", arg))
for arg in cmd_line
os.path.join(self.base_dir or "./", dir) for dir in args.directories
]
directories = [dir for dir in directories if os.path.isdir(dir)]
if len(directories) == 0 and self.base_dir:
directories = [self.base_dir]
mode = ["-m", args.mode] if args.mode else []
num_images = ["-n", args.num_images] if args.num_images else []
with NamedTemporaryFile(suffix=".png") as ntf:
self.logger.debug(
f"Running: "
+ str(
[
self.kollagen_path,
"-r",
*directories,
"-w",
args.width,
"-h",
args.height,
*mode,
"-o",
ntf.name,
]
)
)
subprocess.run(
[self.kollagen_path, "-r", *directories, "-o", ntf.name],
[self.kollagen_path, "-r", *directories, *mode, "-o", ntf.name],
check=True,
capture_output=True,
)
ntf.seek(0)
update.message.reply_photo(ntf)
return True
def tg_list_modes(self, update: Update, context: CallbackContext):
modes = subprocess.run(
[self.kollagen_path, "-m"], check=True, capture_output=True
def tg_help(self, update: Update, context: CallbackContext):
update.message.reply_text(
f"```{self.parser.format_help()}```", parse_mode=ParseMode.MARKDOWN
)
update.message.reply_text(f"Available modes: {modes.stdout.decode('utf-8')}")
def tg_error(self, update: object, context: CallbackContext) -> None:
self.logger.error(
@ -99,7 +180,7 @@ class KollagenBot:
def main() -> None:
logging.basicConfig(
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
level=logging.INFO,
level=getattr(logging, os.getenv("LOG_LEVEL", "info").upper()),
)
tg_token = os.getenv("TG_TOKEN")

16
tgbot/kollagen-bot/parser.py Normal file
View File

@ -0,0 +1,16 @@
import argparse
import re
class ArgumentParserError(RuntimeError):
pass
class SafeArgumentParser(argparse.ArgumentParser):
def error(self, message):
raise ArgumentParserError(message)
def safe_str(val: str):
if re.findall(r'[^\w]', val):
raise RuntimeError("No special characters in arguments allowed!")
return val