mirror of
https://github.com/myned/modufur.git
synced 2024-11-01 13:02:38 +00:00
Change to black formatting
This commit is contained in:
Myned
parent
5956d9301e
commit
2c171f204c
7 changed files with 117 additions and 98 deletions
|
|
@ -6,70 +6,74 @@ import pysaucenao
|
|||
from tools import components, scraper
|
||||
|
||||
|
||||
plugin = lightbulb.Plugin('booru')
|
||||
plugin = lightbulb.Plugin("booru")
|
||||
extractor = urlextract.URLExtract()
|
||||
|
||||
|
||||
@plugin.command
|
||||
#@lightbulb.option('attachment', 'Attachment(s) to reverse', required=False)
|
||||
@lightbulb.option('url', 'URL(s) to reverse, separated by space')
|
||||
@lightbulb.command('reverse', 'Reverse image search using SauceNAO & Kheina', ephemeral=True)
|
||||
# @lightbulb.option('attachment', 'Attachment(s) to reverse')
|
||||
@lightbulb.option("url", "URL(s) to reverse, separated by space")
|
||||
@lightbulb.command("reverse", "Reverse image search using SauceNAO & Kheina", ephemeral=True)
|
||||
@lightbulb.implements(lightbulb.SlashCommand, lightbulb.MessageCommand)
|
||||
async def reverse(context):
|
||||
match context:
|
||||
case lightbulb.SlashContext():
|
||||
urls = extractor.find_urls(context.options.url or '', only_unique=True, with_schema_only=True)
|
||||
urls = extractor.find_urls(context.options.url or "", only_unique=True, with_schema_only=True)
|
||||
|
||||
if not urls:
|
||||
await context.respond('**Invalid URL(s).**')
|
||||
await context.respond("**Invalid URL(s).**")
|
||||
return
|
||||
|
||||
await _reverse(context, urls)
|
||||
case lightbulb.MessageContext():
|
||||
urls = extractor.find_urls(context.options.target.content or '', only_unique=True, with_schema_only=True)
|
||||
urls = extractor.find_urls(context.options.target.content or "", only_unique=True, with_schema_only=True)
|
||||
urls += [attachment.url for attachment in context.options.target.attachments if attachment.url not in urls]
|
||||
|
||||
if not urls:
|
||||
await context.respond('**No images found.**')
|
||||
await context.respond("**No images found.**")
|
||||
return
|
||||
|
||||
selector = None
|
||||
|
||||
if len(urls) > 1:
|
||||
selector = components.Selector(
|
||||
pages=[f'**Select potential images to search: `{urls.index(url) + 1}/{len(urls)}`**\n{url}' for url in urls],
|
||||
pages=[
|
||||
f"**Select potential images to search: `{urls.index(url) + 1}/{len(urls)}`**\n{url}"
|
||||
for url in urls
|
||||
],
|
||||
buttons=[components.Back(), components.Forward(), components.Select(), components.Confirm()],
|
||||
urls=urls
|
||||
urls=urls,
|
||||
)
|
||||
|
||||
await selector.send(context.interaction, ephemeral=True)
|
||||
await selector.wait()
|
||||
|
||||
if selector.timed_out:
|
||||
await context.interaction.edit_initial_response('**Timed out.**', components=None)
|
||||
await context.interaction.edit_initial_response("**Timed out.**", components=None)
|
||||
return
|
||||
|
||||
urls = selector.selected
|
||||
|
||||
await _reverse(context, urls, selector=selector)
|
||||
|
||||
|
||||
@reverse.set_error_handler()
|
||||
async def on_reverse_error(event):
|
||||
error = None
|
||||
|
||||
match event.exception.__cause__:
|
||||
case pysaucenao.ShortLimitReachedException():
|
||||
error = '**API limit reached. Please try again in a minute.**'
|
||||
error = "**API limit reached. Please try again in a minute.**"
|
||||
case pysaucenao.DailyLimitReachedException():
|
||||
error = '**Daily API limit reached. Please try again tomorrow.**'
|
||||
error = "**Daily API limit reached. Please try again tomorrow.**"
|
||||
case pysaucenao.FileSizeLimitException() as url:
|
||||
error = f'**Image file size too large:**\n{url}'
|
||||
error = f"**Image file size too large:**\n{url}"
|
||||
case pysaucenao.ImageSizeException() as url:
|
||||
error = f'**Image resolution too small:**\n{url}'
|
||||
error = f"**Image resolution too small:**\n{url}"
|
||||
case pysaucenao.InvalidImageException() as url:
|
||||
error = f'**Invalid image:**\n{url}'
|
||||
error = f"**Invalid image:**\n{url}"
|
||||
case pysaucenao.UnknownStatusCodeException():
|
||||
error = '**An unknown SauceNAO error has occurred. The service may be down.**'
|
||||
error = "**An unknown SauceNAO error has occurred. The service may be down.**"
|
||||
|
||||
if error:
|
||||
try:
|
||||
|
|
@ -79,6 +83,7 @@ async def on_reverse_error(event):
|
|||
|
||||
return True
|
||||
|
||||
|
||||
async def _reverse(context, urls, *, selector=None):
|
||||
if not selector:
|
||||
await context.respond(hikari.ResponseType.DEFERRED_MESSAGE_CREATE)
|
||||
|
|
@ -87,24 +92,27 @@ async def _reverse(context, urls, *, selector=None):
|
|||
|
||||
if not matches:
|
||||
if selector:
|
||||
await context.interaction.edit_initial_response('**No matches found.**', components=None)
|
||||
await context.interaction.edit_initial_response("**No matches found.**", components=None)
|
||||
else:
|
||||
await context.respond('**No matches found.**')
|
||||
await context.respond("**No matches found.**")
|
||||
return
|
||||
|
||||
pages = [(hikari.Embed(
|
||||
title=match['artist'], url=match['url'], color=context.get_guild().get_my_member().get_top_role().color)
|
||||
.set_author(name=f'{match["similarity"]}% Match')
|
||||
.set_image(match['thumbnail'])
|
||||
.set_footer(match['source']))
|
||||
if match else f'**No match found.**\n{urls[index]}' for index, match in enumerate(matches)]
|
||||
pages = [
|
||||
(
|
||||
hikari.Embed(
|
||||
title=match["artist"], url=match["url"], color=context.get_guild().get_my_member().get_top_role().color
|
||||
)
|
||||
.set_author(name=f'{match["similarity"]}% Match')
|
||||
.set_image(match["thumbnail"])
|
||||
.set_footer(match["source"])
|
||||
)
|
||||
if match
|
||||
else f"**No match found.**\n{urls[index]}"
|
||||
for index, match in enumerate(matches)
|
||||
]
|
||||
|
||||
if len(pages) > 1:
|
||||
selector = components.Selector(
|
||||
pages=pages,
|
||||
buttons=[components.Back(), components.Forward()],
|
||||
timeout=900
|
||||
)
|
||||
selector = components.Selector(pages=pages, buttons=[components.Back(), components.Forward()], timeout=900)
|
||||
|
||||
await selector.send_edit(context.interaction)
|
||||
else:
|
||||
|
|
@ -113,7 +121,10 @@ async def _reverse(context, urls, *, selector=None):
|
|||
else:
|
||||
await context.respond(pages[0])
|
||||
|
||||
|
||||
def load(bot):
|
||||
bot.add_plugin(plugin)
|
||||
|
||||
|
||||
def unload(bot):
|
||||
bot.remove_plugin(plugin)
|
||||
|
|
|
|||
|
|
@ -1,32 +1,37 @@
|
|||
import os
|
||||
|
||||
import lightbulb
|
||||
|
||||
|
||||
plugin = lightbulb.Plugin('master')
|
||||
plugin = lightbulb.Plugin("master")
|
||||
|
||||
|
||||
@plugin.command
|
||||
@lightbulb.option('command', 'What is your command, master?', required=False, choices=('reload', 'sleep'))
|
||||
@lightbulb.command('master', 'Commands my master can demand of me', ephemeral=True)
|
||||
@lightbulb.option("command", "What is your command, master?", required=False, choices=("reload", "sleep"))
|
||||
@lightbulb.command("master", "Commands my master can demand of me", ephemeral=True)
|
||||
@lightbulb.implements(lightbulb.SlashCommand)
|
||||
async def master(context):
|
||||
if context.user.id == context.bot.application.owner.id:
|
||||
match context.options.command:
|
||||
case 'reload':
|
||||
case "reload":
|
||||
context.bot.reload_extensions(*context.bot.extensions)
|
||||
|
||||
extensions = [os.path.splitext(extension)[1][1:] for extension in context.bot.extensions]
|
||||
await context.respond(f'**Reloaded `{"`, `".join(extensions[:-1])}`, and `{extensions[-1]}` for you, master.**')
|
||||
case 'sleep':
|
||||
await context.respond('**Goodnight, master.**')
|
||||
await context.respond(
|
||||
f'**Reloaded `{"`, `".join(extensions[:-1])}`, and `{extensions[-1]}` for you, master.**'
|
||||
)
|
||||
case "sleep":
|
||||
await context.respond("**Goodnight, master.**")
|
||||
await context.bot.close()
|
||||
case _:
|
||||
await context.respond(f'**Hello, master.**')
|
||||
await context.respond(f"**Hello, master.**")
|
||||
else:
|
||||
await context.respond(f'**{context.bot.application.owner.mention} is my master. 🐺**')
|
||||
await context.respond(f"**{context.bot.application.owner.mention} is my master. 🐺**")
|
||||
|
||||
|
||||
def load(bot):
|
||||
bot.add_plugin(plugin)
|
||||
|
||||
|
||||
def unload(bot):
|
||||
bot.remove_plugin(plugin)
|
||||
|
|
|
|||
19
config.py
19
config.py
|
|
@ -1,29 +1,32 @@
|
|||
import toml
|
||||
import hikari
|
||||
|
||||
|
||||
ACTIVITY = hikari.ActivityType.LISTENING
|
||||
ERROR = '```❗ An internal error has occurred. This has been reported to my master. 🐺```'
|
||||
CONFIG = '''\
|
||||
ERROR = "```❗ An internal error has occurred. This has been reported to my master. 🐺```"
|
||||
CONFIG = """\
|
||||
guilds = [] # guild IDs to register commands, empty for global
|
||||
client = 0 # bot application ID
|
||||
token = "" # bot token
|
||||
activity = "" # bot status
|
||||
saucenao = "" # saucenao token
|
||||
e621 = "" # e621 token
|
||||
'''
|
||||
"""
|
||||
|
||||
|
||||
try:
|
||||
config = toml.load('config.toml')
|
||||
config = toml.load("config.toml")
|
||||
except FileNotFoundError:
|
||||
with open('config.toml', 'w') as f:
|
||||
with open("config.toml", "w") as f:
|
||||
f.write(CONFIG)
|
||||
print('config.toml created with default values. Restart when modified.')
|
||||
print("config.toml created with default values. Restart when modified.")
|
||||
exit()
|
||||
|
||||
|
||||
def error(event):
|
||||
exception = event.exception.__cause__ or event.exception
|
||||
|
||||
return (f'**`{event.context.command.name}` in {event.context.get_channel().mention}'
|
||||
f'```❗ {type(exception).__name__}: {exception}```**')
|
||||
return (
|
||||
f"**`{event.context.command.name}` in {event.context.get_channel().mention}"
|
||||
f"```❗ {type(exception).__name__}: {exception}```**"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -18,6 +18,10 @@ hikari-miru = "*"
|
|||
pysaucenao = {git = "https://github.com/Myned/pysaucenao.git"}
|
||||
|
||||
[tool.poetry.dev-dependencies]
|
||||
black = "*"
|
||||
|
||||
[tool.black]
|
||||
line-length = 120
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core>=1.0.0"]
|
||||
|
|
|
|||
12
run.py
12
run.py
|
|
@ -1,4 +1,5 @@
|
|||
import os
|
||||
|
||||
import hikari
|
||||
import lightbulb
|
||||
import miru
|
||||
|
|
@ -8,13 +9,12 @@ import config as c
|
|||
|
||||
# Unix optimizations
|
||||
# https://github.com/hikari-py/hikari#uvloop
|
||||
if os.name != 'nt':
|
||||
if os.name != "nt":
|
||||
import uvloop
|
||||
|
||||
uvloop.install()
|
||||
|
||||
bot = lightbulb.BotApp(
|
||||
token=c.config['token'],
|
||||
default_enabled_guilds=c.config['guilds'])
|
||||
bot = lightbulb.BotApp(token=c.config["token"], default_enabled_guilds=c.config["guilds"])
|
||||
|
||||
|
||||
@bot.listen(lightbulb.CommandErrorEvent)
|
||||
|
|
@ -30,5 +30,5 @@ async def on_error(event):
|
|||
|
||||
|
||||
miru.load(bot)
|
||||
bot.load_extensions_from('tools', 'commands')
|
||||
bot.run(activity=hikari.Activity(name=c.config['activity'], type=c.ACTIVITY))
|
||||
bot.load_extensions_from("tools", "commands")
|
||||
bot.run(activity=hikari.Activity(name=c.config["activity"], type=c.ACTIVITY))
|
||||
|
|
|
|||
|
|
@ -3,48 +3,35 @@ import lightbulb
|
|||
from miru.ext import nav
|
||||
|
||||
|
||||
plugin = lightbulb.Plugin('components')
|
||||
plugin = lightbulb.Plugin("components")
|
||||
|
||||
|
||||
class Back(nav.PrevButton):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
style=hikari.ButtonStyle.SECONDARY,
|
||||
label='⟵',
|
||||
emoji=None
|
||||
)
|
||||
super().__init__(style=hikari.ButtonStyle.SECONDARY, label="⟵", emoji=None)
|
||||
|
||||
|
||||
class Forward(nav.NextButton):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
style=hikari.ButtonStyle.SECONDARY,
|
||||
label='⟶',
|
||||
emoji=None
|
||||
)
|
||||
super().__init__(style=hikari.ButtonStyle.SECONDARY, label="⟶", emoji=None)
|
||||
|
||||
|
||||
class Confirm(nav.StopButton):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
style=hikari.ButtonStyle.PRIMARY,
|
||||
label='➤',
|
||||
emoji=None
|
||||
)
|
||||
super().__init__(style=hikari.ButtonStyle.PRIMARY, label="➤", emoji=None)
|
||||
|
||||
async def callback(self, context):
|
||||
await context.edit_response(content='**Searching...**', components=None)
|
||||
await context.edit_response(content="**Searching...**", components=None)
|
||||
|
||||
self.view.stop()
|
||||
|
||||
async def before_page_change(self):
|
||||
self.disabled = False if self.view.selected else True
|
||||
|
||||
|
||||
class Select(nav.NavButton):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
style=hikari.ButtonStyle.DANGER,
|
||||
label='✗',
|
||||
emoji=None
|
||||
)
|
||||
super().__init__(style=hikari.ButtonStyle.DANGER, label="✗", emoji=None)
|
||||
|
||||
async def callback(self, context):
|
||||
if self.view.urls[self.view.current_page] not in self.view.selected:
|
||||
|
|
@ -64,7 +51,7 @@ class Select(nav.NavButton):
|
|||
|
||||
def _button(self, *, selected=False):
|
||||
self.style = hikari.ButtonStyle.SUCCESS if selected else hikari.ButtonStyle.DANGER
|
||||
self.label = '✔' if selected else '✗'
|
||||
self.label = "✔" if selected else "✗"
|
||||
|
||||
try:
|
||||
confirm = next((child for child in self.view.children if isinstance(child, Confirm)))
|
||||
|
|
@ -75,11 +62,7 @@ class Select(nav.NavButton):
|
|||
|
||||
class Selector(nav.NavigatorView):
|
||||
def __init__(self, *, pages=[], buttons=[], timeout=120, urls=[]):
|
||||
super().__init__(
|
||||
pages=pages,
|
||||
buttons=buttons,
|
||||
timeout=timeout
|
||||
)
|
||||
super().__init__(pages=pages, buttons=buttons, timeout=timeout)
|
||||
self.urls = urls
|
||||
self.selected = []
|
||||
self.saved = set()
|
||||
|
|
@ -110,5 +93,7 @@ class Selector(nav.NavigatorView):
|
|||
|
||||
def load(bot):
|
||||
bot.add_plugin(plugin)
|
||||
|
||||
|
||||
def unload(bot):
|
||||
bot.remove_plugin(plugin)
|
||||
|
|
|
|||
|
|
@ -6,13 +6,14 @@ import pysaucenao
|
|||
import config as c
|
||||
|
||||
|
||||
plugin = lightbulb.Plugin('scraper')
|
||||
sauce = pysaucenao.SauceNao(api_key=c.config['saucenao'], priority=(29, 40, 41)) # e621 > Fur Affinity > Twitter
|
||||
plugin = lightbulb.Plugin("scraper")
|
||||
sauce = pysaucenao.SauceNao(api_key=c.config["saucenao"], priority=(29, 40, 41)) # e621 > Fur Affinity > Twitter
|
||||
|
||||
|
||||
async def reverse(urls):
|
||||
return [await _saucenao(url) or await _kheina(url) for url in urls]
|
||||
|
||||
|
||||
async def _saucenao(url):
|
||||
try:
|
||||
results = await sauce.from_url(url)
|
||||
|
|
@ -23,35 +24,45 @@ async def _saucenao(url):
|
|||
except pysaucenao.InvalidImageException:
|
||||
raise pysaucenao.InvalidImageException(url)
|
||||
|
||||
return {
|
||||
'url': results[0].url,
|
||||
'artist': ', '.join(results[0].authors) or 'Unknown',
|
||||
'thumbnail': results[0].thumbnail,
|
||||
'similarity': round(results[0].similarity),
|
||||
'source': tldextract.extract(results[0].index).domain
|
||||
} if results else None
|
||||
return (
|
||||
{
|
||||
"url": results[0].url,
|
||||
"artist": ", ".join(results[0].authors) or "Unknown",
|
||||
"thumbnail": results[0].thumbnail,
|
||||
"similarity": round(results[0].similarity),
|
||||
"source": tldextract.extract(results[0].index).domain,
|
||||
}
|
||||
if results
|
||||
else None
|
||||
)
|
||||
|
||||
|
||||
async def _kheina(url):
|
||||
content = await _post('https://api.kheina.com/v1/search', {'url': url})
|
||||
content = await _post("https://api.kheina.com/v1/search", {"url": url})
|
||||
|
||||
if content['results'][0]['similarity'] < 50:
|
||||
if content["results"][0]["similarity"] < 50:
|
||||
return None
|
||||
|
||||
return {
|
||||
'url': content['results'][0]['sources'][0]['source'],
|
||||
'artist': content['results'][0]['sources'][0]['artist'] or 'Unknown',
|
||||
'thumbnail': f'https://cdn.kheina.com/file/kheinacom/{content["results"][0]["sources"][0]["sha1"]}.jpg',
|
||||
'similarity': round(content['results'][0]['similarity']),
|
||||
'source': tldextract.extract(content['results'][0]['sources'][0]['source']).domain
|
||||
"url": content["results"][0]["sources"][0]["source"],
|
||||
"artist": content["results"][0]["sources"][0]["artist"] or "Unknown",
|
||||
"thumbnail": f'https://cdn.kheina.com/file/kheinacom/{content["results"][0]["sources"][0]["sha1"]}.jpg',
|
||||
"similarity": round(content["results"][0]["similarity"]),
|
||||
"source": tldextract.extract(content["results"][0]["sources"][0]["source"]).domain,
|
||||
}
|
||||
|
||||
|
||||
async def _post(url, data):
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, data=data, headers={'User-Agent': 'Myned/Modufur (https://github.com/Myned/Modufur)'}) as response:
|
||||
async with session.post(
|
||||
url, data=data, headers={"User-Agent": "Myned/Modufur (https://github.com/Myned/Modufur)"}
|
||||
) as response:
|
||||
return await response.json() if response.status == 200 else None
|
||||
|
||||
|
||||
def load(bot):
|
||||
bot.add_plugin(plugin)
|
||||
|
||||
|
||||
def unload(bot):
|
||||
bot.remove_plugin(plugin)
|
||||
|
|
|
|||
Loading…
Reference in a new issue