add argparser for booru command

This commit is contained in:
Owl 2025-01-05 07:56:09 +07:00
parent f8051314cc
commit c65401abc7
3 changed files with 168 additions and 68 deletions

View File

@ -2,4 +2,6 @@ import os
TOKEN = os.getenv("TOKEN")
BOORU_API_URL = "http://booru-api:3456/booru"
INLINE_QUERY_CACHE_SECONDS = 30
STICKER_SHITHOLE = -1002471390283

View File

@ -1,10 +1,19 @@
import requests
from telegram import Update, InlineQueryResultPhoto
from telegram import Update
from telegram.ext import ContextTypes
from tgbot.config import STICKER_SHITHOLE, BOORU_API_URL
from tgbot.shit.hentai import QueryParams
from tgbot.config import STICKER_SHITHOLE, BOORU_API_URL, INLINE_QUERY_CACHE_SECONDS
from tgbot.shit.hentai import parse_args, create_parser
from tgbot.shit.render import render_text_on_image
import re
def escape_markdown(text):
"""
Escapes special characters for MarkdownV2 formatting.
"""
escape_chars = r"[_*[\]()~`>#+-=|{}.!]"
return re.sub(f"([{escape_chars}])", r"\\\1", text)
async def handle_xitter(update: Update, context: ContextTypes.DEFAULT_TYPE):
@ -44,30 +53,62 @@ async def handle_red_ebalo(update: Update, context: ContextTypes.DEFAULT_TYPE):
async def handle_hentai(update: Update, context: ContextTypes.DEFAULT_TYPE):
query = update.inline_query.query.replace("hentai", "").replace("", "--")
parser = create_parser()
try:
args = parse_args(query, parser)
except RuntimeError as e:
if getattr(parser, 'help_message', None):
help_message = parser.help_message
escaped_help_message = escape_markdown(help_message)
help_response = [
{
"type": "article",
"id": "1",
"title": "Help message",
"description": "Print help",
"input_message_content": {
"message_text": f"```\n{escaped_help_message}\n```",
"parse_mode": "MarkdownV2",
}
}
]
return await context.bot.answer_inline_query(update.inline_query.id, help_response,
cache_time=INLINE_QUERY_CACHE_SECONDS)
else:
error_message = str(e)
escaped_error_message = escape_markdown(error_message)
error_response = [
{
"type": "article",
"id": "1",
"title": "ERROR OCCURRED! That means, at least one of us is a retard.",
"description": "None",
"title": "Invalid command!",
"description": error_message,
"input_message_content": {
"message_text": "None"
"message_text": f"```\n{escaped_error_message}\n```",
"parse_mode": "MarkdownV2",
}
}
]
query = update.inline_query.query
return await context.bot.answer_inline_query(update.inline_query.id, error_response,
cache_time=INLINE_QUERY_CACHE_SECONDS)
try:
params = QueryParams.from_query_str(query).__dict__
except ValueError as e:
print(e)
error_response[0]["description"] = "error parsing query"
error_response[0]["input_message_content"]["message_text"] = "error parsing query"
return await context.bot.answer_inline_query(update.inline_query.id, error_response)
print(args)
response = requests.post(BOORU_API_URL, json=params)
response = requests.post(BOORU_API_URL, json=args)
if response.status_code == 200:
@ -79,24 +120,51 @@ async def handle_hentai(update: Update, context: ContextTypes.DEFAULT_TYPE):
booru_post_url = post["booru"]["domain"] + \
post["booru"]["site"]["api"]["postView"] + post["id"]
results.append(
InlineQueryResultPhoto(
id="%d" % idx,
photo_url=post["fileUrl"],
thumbnail_url=post["previewUrl"],
title="IMAGE #%d" % idx,
description="Image from booru",
caption="source -> " + booru_post_url
)
)
if not args["spoiler"]:
await context.bot.answer_inline_query(update.inline_query.id, results)
results.append(
{
"type": "photo",
"id": "%d" % idx,
"photo_url": post["fileUrl"],
"thumb_url": post["previewUrl"],
"caption": "source -> " + booru_post_url
}
)
else:
print(response.text)
results.append({
"type": "photo",
"id": "%d" % idx,
"photo_url": post["fileUrl"],
"thumb_url": post["previewUrl"],
"caption": "source -> " + booru_post_url,
"input_message_content": {
"message_text": "image with tags = " + str(args["tags"]) + " -> " + booru_post_url,
"has_spoiler": True
}
})
await context.bot.answer_inline_query(update.inline_query.id, results,
cache_time=INLINE_QUERY_CACHE_SECONDS)
else:
error_response = [
{
"type": "article",
"id": "1",
"title": "ERROR!",
"description": "None",
"input_message_content": {
"message_text": "None"
}
}
]
error_response[0]["description"] = "booru error"
error_response[0]["input_message_content"]["message_text"] = response.text
return await context.bot.answer_inline_query(update.inline_query.id, error_response)
return await context.bot.answer_inline_query(update.inline_query.id, error_response,
cache_time=INLINE_QUERY_CACHE_SECONDS)

View File

@ -1,46 +1,76 @@
import re
import argparse
import io
import shlex
def combine_whitespaces(s: str) -> str:
return re.compile(r"\s+").sub(" ", s).strip()
class CustomArgParser(argparse.ArgumentParser):
def __init__(
self,
*args, **kwargs,
):
super().__init__(*args, **kwargs)
self.help_message = None
def print_help(self, file=None):
# Store the help message in a buffer instead of printing
if file is None:
help_buffer = io.StringIO()
super().print_help(file=help_buffer)
self.help_message = help_buffer.getvalue()
help_buffer.close()
else:
super().print_help(file=file)
def parse_args(self, args=None, namespace=None):
# Check for --help manually to avoid stdout output
if '--help' in args:
self.print_help()
raise RuntimeError("Help requested")
return super().parse_args(args, namespace)
def exit(self, status=0, message=None):
raise RuntimeError(message)
def error(self, message):
raise RuntimeError(f"Error: {message}")
class QueryParams:
def __init__(self, booru: str, tags: list[str], random: bool = True, limit: int = 10):
self.booru = booru
self.tags = tags
self.random = random
self.limit = limit
def create_parser():
parser = CustomArgParser(
description="A command to search random stuff on various image boorus. Feel free to use it irresponsibly!",
usage="@owlrandomshitbot hentai BOORU_NAME --tags TAG1,TAG2,TAG3 -l LIMIT --random -s",
)
@classmethod
def from_query_str(cls, query: str):
query = query.replace("hentai ", "")
parser.add_argument("booru", type=str,
choices=['e6', 'e621', 'e9', 'e926', 'hh', 'hypno', 'hypnohub', 'db', 'dan', 'danbooru', 'kc',
'konac', 'kcom', 'kn', 'konan', 'knet', 'yd', 'yand', 'yandere', 'gb', 'gel',
'gelbooru',
'r34', 'rule34', 'sb', 'safe', 'safebooru', 'tb', 'tbib', 'big', 'xb', 'xbooru', 'pa',
'paheal', 'dp', 'derp', 'derpi', 'derpibooru', 'rb', 'realbooru'],
help="Booru to search. For context, this command uses this library"
" -> https://github.com/AtoraSuunva/booru"
)
query = combine_whitespaces(query)
parser.add_argument("--tags", "-t",
dest="tags", type=str, help="Tags, comma separated. If query contains tags with spaces, "
"e.g. 'downvote bait', put it in quotes.")
list_ = query.split(" ")
parser.add_argument("--random", "-r",
dest="random", action='store_true', default=False,
help="Randomize the output, default=False")
booru = list_[0]
parser.add_argument("--limit", "-l",
dest="limit", type=int, default=20, help="Limit the number of output posts, default=20")
list_.remove(booru)
parser.add_argument("--spoiler", "-s", dest="spoiler",
action="store_true", help="Send the image with a spoiler.")
tags = []
random = False
limit = 10
return parser
for e in list_:
if "random=" in e: # this is 'random' argument
random = e.replace("random=", "").lower() == "true"
continue
if "limit=" in e: # this is 'limit' argument
limit = int(e.replace("limit=", ""))
continue
def parse_args(args: str, parser: CustomArgParser):
args = parser.parse_args(shlex.split(args)).__dict__
tags.append(e)
args["tags"] = args["tags"].split(",")
return cls(booru, tags, random, limit)
def __str__(self):
return "QueryParams(booru=" + self.booru + ", tags=" + str(self.tags) + ", random=" + str(self.random) + \
", limit=" + str(self.limit) + ")"
return args