microbot/ubot/modules/booru.py

from random import choice
from re import compile

from asyncbooru import Danbooru, Gelbooru, Konachan, Sankaku, Yandere
from telethon import Button

from ubot import ldr

help_str = "Fetches images from Danbooru, Gelbooru, Konachan, Sankaku Complex and Yandere, takes tags as arguments."
ext_regex = compile(r"\.(jpg|jpeg|png|mp4|gif)($|\?)")

dan_api = Danbooru(ldr.aioclient)
gel_api = Gelbooru(ldr.aioclient)
kon_api = Konachan(ldr.aioclient)
san_api = Sankaku(ldr.aioclient)
yan_api = Yandere(ldr.aioclient)

dan_butt = {}
gel_butt = {}
kon_butt = {}
san_butt = {}
yan_butt = {}

normal_commands = {
    "dan": dan_api,
    "gel": gel_api,
    "kon": kon_api,
    "san": san_api,
    "yan": yan_api
}

button_commands = {
    "danb": [dan_api, dan_butt, "dan"],
    "gelb": [gel_api, gel_butt, "gel"],
    "konb": [kon_api, kon_butt, "kon"],
    "sanb": [san_api, san_butt, "san"],
    "yanb": [yan_api, yan_butt, "yan"]
}


@ldr.add_dict(normal_commands, pattern_extra="(s)(f|)", help=help_str, userlocking=True)
@ldr.add_dict(normal_commands, pattern_extra="(x|q|)(f|)", userlocking=True, nsfw=True, nsfw_warning="NSFW commands are disabled in this chat, add 's' to the end of the command for SFW images.", hide_help=True)
async def booru(event):
    posts = await event.extra.get_random_posts(event.args, 3, event.other_args[0])

    if not posts:
        await event.reply(f"No results for query: {event.args}")
        return

    images = [[post.file_url, post.sauce, post.source, post.rating] for post in posts if post.file_url and ext_regex.search(post.file_url)]

    if not images:
        await event.reply(f"Failed to find URLs for query: {event.args}")
        return

    for image in images:
        try:
            await event.reply(
                gen_source_string(image[1], image[2]),
                file=image[0],
                force_document=bool(event.other_args[1]),
                spoiler=event.chat_db.spoiler_nsfw and event.extra._get_rating("x") == image[3]
            )
            return
        except:
            pass

    await event.reply(f"Failed to fetch media for query: {event.args}")


@ldr.add_inline_photo("dan(s|x|q|)", default="dan", extra=dan_api)
@ldr.add_inline_photo("gel(s|x|q|)", default="gel", extra=gel_api)
@ldr.add_inline_photo("kon(s|x|q|)", default="kon", extra=kon_api)
@ldr.add_inline_photo("san(s|x|q|)", default="san", extra=san_api)
@ldr.add_inline_photo("yan(s|x|q|)", default="yan", extra=yan_api)
async def booru_inline(event):
    posts = await event.extra.get_random_posts(event.args, 3, event.other_args[0])
    return [[post.file_url, gen_source_string(post.sauce, post.source)] for post in posts if post.file_url and ext_regex.search(post.file_url)] if posts else None


@ldr.add_dict(button_commands, pattern_extra="(s)", help=help_str, userlocking=True)
@ldr.add_dict(button_commands, pattern_extra="(x|q|)", userlocking=True, nsfw=True, nsfw_warning="NSFW commands are disabled in this chat, add 's' to the end of the command for SFW images.", hide_help=True)
async def booru_buttons(event):
    posts = await event.extra[0].get_random_posts(event.args, 30, event.other_args[0])

    if not posts:
        await event.reply(f"No results for query: {event.args}")
        return

    images = [[post.file_url, post.sauce, post.source, post.rating] for post in posts if post.file_url and ext_regex.search(post.file_url)]

    if not images:
        await event.reply(f"Failed to find URLs for query: {event.args}")
        return

    event.extra[1][f"{event.chat.id}_{event.id}"] = [0, images]

    await event.reply(
        gen_source_string(images[0][1], images[0][2]),
        file=images[0][0],
        buttons=[Button.inline('➡️', f'{event.extra[2]}*{event.chat.id}_{event.id}*r')],
        spoiler=event.chat_db.spoiler_nsfw and event.extra[0]._get_rating("x") == images[0][3]
    )


@ldr.add_callback_query("dan", extra=[dan_butt, dan_api])
@ldr.add_callback_query("gel", extra=[gel_butt, gel_api])
@ldr.add_callback_query("kon", extra=[kon_butt, kon_api])
@ldr.add_callback_query("san", extra=[san_butt, san_api])
@ldr.add_callback_query("yan", extra=[yan_butt, yan_api])
async def booru_buttons_callback(event):
    args_split = event.args.split("*")

    dict_id = args_split[0]
    direction = args_split[1]
    this_image = None

    if dict_id in event.extra[0]:
        this_dict = event.extra[0][dict_id]
    else:
        await event.answer("The bot was restarted and the buttons on this message can no longer be used.", alert=True)
        return

    if direction == "r":
        this_dict[0] += 1

        if this_dict[0] + 1 > len(this_dict[1]):
            this_dict[0] = len(this_dict[1]) - 1

        this_image = this_dict[1][this_dict[0]]
    elif direction == "l":
        this_dict[0] -= 1

        if this_dict[0] < 0:
            this_dict[0] = 0

        this_image = this_dict[1][this_dict[0]]

    buttons = []

    if this_dict[0] > 0:
        buttons += [Button.inline('⬅️', f'{event.command}*{dict_id}*l')]

    if len(this_dict[1]) - 1 > this_dict[0]:
        buttons += [Button.inline('➡️', f'{event.command}*{dict_id}*r')]

    try:
        await event.edit(
            gen_source_string(this_image[1], this_image[2]),
            file=this_image[0],
            buttons=buttons,
            spoiler=event.chat_db.spoiler_nsfw and event.extra[1]._get_rating("x") == this_image[3]
        )
    except:
        await event.answer(
            choice(
                [
                    "I shidded myself on that one.",
                    "Oopsie poopsies.",
                    "Something went wrong, I don't know what it was though.",
                    "Yikes, something broke!",
                    "Nick can't code and something went wrong.",
                    "Well, this is embarrassing.",
                    "I tripped and fell trying to handle that button press.",
                    "I've failed you for the last time! *dies*",
                    "That shouldn't happen.",
                    "Try again maybe? :(",
                    "I broke.",
                    "I've fallen and I can't handle that request!"
                ]
            )
        )


def gen_source_string(source: str, orig_source: str) -> str:
    if orig_source and orig_source.startswith("http"):
        return f"[source]({source}) - [original source]({orig_source.split()[0]})"
    else:
        return f"[source]({source})"