77 lines
2.9 KiB
Python
77 lines
2.9 KiB
Python
import argparse
|
|
import io
|
|
import shlex
|
|
|
|
|
|
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}")
|
|
|
|
|
|
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",
|
|
)
|
|
|
|
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"
|
|
)
|
|
|
|
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.")
|
|
|
|
parser.add_argument("--random", "-r",
|
|
dest="random", action='store_true', default=False,
|
|
help="Randomize the output, default=False")
|
|
|
|
parser.add_argument("--limit", "-l",
|
|
dest="limit", type=int, default=20, help="Limit the number of output posts, default=20")
|
|
|
|
parser.add_argument("--spoiler", "-s", dest="spoiler",
|
|
action="store_true", help="Send the image with a spoiler.")
|
|
|
|
return parser
|
|
|
|
|
|
def parse_args(args: str, parser: CustomArgParser):
|
|
args = parser.parse_args(shlex.split(args)).__dict__
|
|
|
|
args["tags"] = args["tags"].split(",")
|
|
|
|
return args
|