add argparser for booru command

tgbot/shit/hentai.py

@@ -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