tgbots/TeleDiffusionBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Decorator error handling

Browse Source
This commit is contained in:
BarsTiger
2023-03-06 23:45:02 +02:00
parent a5ce82f148
commit ebdd27ece8
7 changed files with 104 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
bot/callbacks/set_model.py
View File

@@ -4,35 +4,18 @@ from aiogram import types
from .factories.set_model import set_model, set_model_page
from bot.keyboards.set_model import get_set_model_keyboard
from bot.utils.private_keyboard import other_user
from bot.keyboards.exception import get_exception_keyboard
from bot.utils.trace_exception import PrettyException
from aiohttp import ClientConnectorError
from bot.modules.api import models
from bot.utils.errorable_command import wrap_exception
@wrap_exception()
async def on_set_model(call: types.CallbackQuery, callback_data: dict):
n = int(callback_data['n'])
temp_message = await call.message.answer('⏳ Setting model...')
try:
await call.message.edit_reply_markup(reply_markup=None)
await models.set_model(db[DBTables.config]['models'][n])
except ClientConnectorError:
await call.answer('❌ Error! Maybe, StableDiffusion API endpoint is incorrect '
'or turned off', show_alert=True)
await call.message.answer('✅ Model set for all users!')
await call.message.delete()
await temp_message.delete()
return
except Exception as e:
exception_id = f'{call.message.message_thread_id}-{call.message.message_id}'
db[DBTables.exceptions][exception_id] = PrettyException(e)
await call.message.reply('❌ Error happened while processing your request', parse_mode='html',
reply_markup=get_exception_keyboard(exception_id))
db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
return
await temp_message.answer('✅ Model set for all users!')
await temp_message.delete()
async def on_page_change(call: types.CallbackQuery, callback_data: dict):

11
bot/handlers/admin/tools.py
View File

@@ -1,11 +1,11 @@
from aiogram import types
from bot.db import db, DBTables
from bot.utils.cooldown import throttle
from bot.keyboards.exception import get_exception_keyboard
from bot.utils.trace_exception import PrettyException
from bot.utils.errorable_command import wrap_exception
from bot.modules.get_hash.get_hash import get_hash
@wrap_exception([IndexError])
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def hash_command(message: types.Message):
try:
@@ -23,10 +23,3 @@ async def hash_command(message: types.Message):
except IndexError:
await message.reply('❌ Reply with this command on PICTURE OR FILE', parse_mode='html')
except Exception as e:
exception_id = f'{message.message_thread_id}-{message.message_id}'
db[DBTables.exceptions][exception_id] = PrettyException(e)
await message.reply('❌ Error happened while processing your request', parse_mode='html',
reply_markup=get_exception_keyboard(exception_id))
return

11
bot/handlers/image_info/image_info.py
View File

@@ -1,11 +1,11 @@
from aiogram import types
from bot.db import db, DBTables
from bot.utils.cooldown import throttle
from bot.keyboards.exception import get_exception_keyboard
from bot.keyboards.image_info import get_img_info_keyboard
from bot.utils.trace_exception import PrettyException
from bot.utils.errorable_command import wrap_exception
@wrap_exception([IndexError])
@throttle(cooldown=10, admin_ids=db[DBTables.config].get('admins'))
async def imginfo(message: types.Message):
try:
@@ -25,10 +25,3 @@ async def imginfo(message: types.Message):
except IndexError:
await message.reply('❌ Reply with this command on PICTURE', parse_mode='html')
except Exception as e:
exception_id = f'{message.message_thread_id}-{message.message_id}'
db[DBTables.exceptions][exception_id] = PrettyException(e)
await message.reply('❌ Error happened while processing your request', parse_mode='html',
reply_markup=get_exception_keyboard(exception_id))
return

22
bot/handlers/txt2img/set_model.py
View File

@@ -2,36 +2,18 @@ from aiogram import types
from bot.db import db, DBTables
from bot.utils.cooldown import throttle
from bot.keyboards.set_model import get_set_model_keyboard
from bot.keyboards.exception import get_exception_keyboard
from bot.utils.trace_exception import PrettyException
from bot.modules.api.models import get_models
from aiohttp import ClientConnectorError
from bot.utils.errorable_command import wrap_exception
@wrap_exception()
@throttle(cooldown=60*60, admin_ids=db[DBTables.config].get('admins'), by_id=False)
async def set_model_command(message: types.Message):
temp_message = await message.reply('⏳ Getting models...')
try:
models = await get_models()
if models is not None and len(models) > 0:
db[DBTables.config]['models'] = models
else:
await temp_message.delete()
await message.reply('❌ No models available')
return
await temp_message.delete()
await message.reply("🪄 You can choose model from available:", reply_markup=get_set_model_keyboard(0))
except ClientConnectorError:
await message.reply('❌ Error! Maybe, StableDiffusion API endpoint is incorrect '
'or turned off')
await temp_message.delete()
return
except Exception as e:
exception_id = f'{message.message_thread_id}-{message.message_id}'
db[DBTables.exceptions][exception_id] = PrettyException(e)
await message.reply('❌ Error happened while processing your request', parse_mode='html',
reply_markup=get_exception_keyboard(exception_id))
return

31
bot/handlers/txt2img/set_settings.py
View File

@@ -2,17 +2,16 @@ from aiogram import types
from bot.db import db, DBTables
from bot.utils.cooldown import throttle
from bot.modules.api.objects.prompt_request import Prompt
from bot.keyboards.exception import get_exception_keyboard
from bot.utils.trace_exception import PrettyException
from bot.utils.errorable_command import wrap_exception
@wrap_exception(custom_loading=True)
async def _set_property(message: types.Message, prop: str, value=None):
temp_message = await message.reply(f"⏳ Setting {prop}...")
if not message.get_args():
await temp_message.edit_text("😶‍🌫️ Specify arguments for this command. Check /help")
return
try:
prompt: Prompt = db[DBTables.prompts].get(message.from_id)
if prompt is None and prop != 'prompt':
await temp_message.edit_text(f"You didn't created any prompt. Specify prompt text at least first time. "
@@ -32,15 +31,6 @@ async def _set_property(message: types.Message, prop: str, value=None):
await message.reply(f'✅ {prop} set')
await temp_message.delete()
except Exception as e:
exception_id = f'{message.message_thread_id}-{message.message_id}'
db[DBTables.exceptions][exception_id] = PrettyException(e)
await message.reply('❌ Error happened while processing your request', parse_mode='html',
reply_markup=get_exception_keyboard(exception_id))
await temp_message.delete()
db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
return
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_prompt_command(message: types.Message):
@@ -125,31 +115,16 @@ async def set_restore_faces_command(message: types.Message):
await _set_property(message, 'restore_faces')
@wrap_exception()
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_sampler_command(message: types.Message):
temp_message = await message.reply('⏳ Getting samplers...')
from bot.modules.api.samplers import get_samplers
from aiohttp import ClientConnectorError
try:
if message.get_args() not in (samplers := await get_samplers()):
await message.reply(
f'❌ You can use only {", ".join(f"<code>{x}</code>" for x in samplers)}',
parse_mode='HTML'
)
return
await temp_message.delete()
except ClientConnectorError:
await message.reply('❌ Error! Maybe, StableDiffusion API endpoint is incorrect '
'or turned off')
await temp_message.delete()
return
except Exception as e:
exception_id = f'{message.message_thread_id}-{message.message_id}'
db[DBTables.exceptions][exception_id] = PrettyException(e)
await message.reply('❌ Error happened while processing your request', parse_mode='html',
reply_markup=get_exception_keyboard(exception_id))
await temp_message.delete()
return
await _set_property(message, 'sampler')

23
bot/handlers/txt2img/txt2img.py
View File

@@ -6,12 +6,11 @@ from bot.modules.api.txt2img import txt2img
from bot.modules.api.objects.get_prompt import get_prompt
from bot.modules.api.objects.prompt_request import Generated
from bot.modules.api.status import wait_for_status
from bot.keyboards.exception import get_exception_keyboard
from bot.keyboards.image_info import get_img_info_keyboard
from bot.utils.trace_exception import PrettyException
from aiohttp import ClientConnectorError
from bot.utils.errorable_command import wrap_exception
@wrap_exception([ValueError], custom_loading=True)
@throttle(cooldown=30, admin_ids=db[DBTables.config].get('admins'))
async def generate_command(message: types.Message):
temp_message = await message.reply("⏳ Enqueued...")
@@ -33,10 +32,10 @@ async def generate_command(message: types.Message):
await wait_for_status()
await temp_message.edit_text(f"⌛ Generating...")
db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
image = await txt2img(prompt)
image_message = await message.reply_photo(photo=image[0])
db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
db[DBTables.generated][image_message.photo[0].file_unique_id] = Generated(
prompt=prompt,
seed=image[1]['seed'],
@@ -50,24 +49,8 @@ async def generate_command(message: types.Message):
await db[DBTables.config].write()
except ClientConnectorError:
await message.reply('❌ Error! Maybe, StableDiffusion API endpoint is incorrect '
'or turned off')
await temp_message.delete()
db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
return
except ValueError as e:
await message.reply(f'❌ Error! {e.args[0]}')
await temp_message.delete()
db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
return
except Exception as e:
exception_id = f'{message.message_thread_id}-{message.message_id}'
db[DBTables.exceptions][exception_id] = PrettyException(e)
await message.reply('❌ Error happened while processing your request', parse_mode='html',
reply_markup=get_exception_keyboard(exception_id))
await temp_message.delete()
db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
return

58
bot/utils/errorable_command.py Normal file
View File

@@ -0,0 +1,58 @@
from functools import wraps
from bot.db import db, DBTables
from bot.keyboards.exception import get_exception_keyboard
from bot.utils.trace_exception import PrettyException
from aiohttp import ClientConnectorError
from aiogram import types
def wrap_exception(unhandled_types: list = None, custom_loading: bool = False):
def decorator(func):
@wraps(func)
async def wrapper(*args, **kwargs):
message = args[0]
if not custom_loading:
if isinstance(message, types.Message):
temp_message = await message.reply('⏳ Loading...')
elif isinstance(message, types.CallbackQuery):
temp_message = await message.message.answer('⏳ Loading...')
else:
raise AttributeError("This wrapper is only for commands!")
try:
_ = await func(*args, **kwargs)
if not custom_loading:
await temp_message.delete()
return _
except ClientConnectorError:
r_string = '❌ Error! Maybe, StableDiffusion API endpoint is incorrect ' \
'or turned off'
if isinstance(message, types.Message):
await message.reply(r_string)
elif isinstance(message, types.CallbackQuery):
await message.message.answer(r_string)
if not custom_loading:
await temp_message.delete()
return
except Exception as e:
if not unhandled_types or e.__class__ not in unhandled_types:
exception_id = f'{message.message_thread_id}-{message.message_id}'
db[DBTables.exceptions][exception_id] = PrettyException(e)
if not custom_loading:
await temp_message.delete()
if isinstance(message, types.Message):
await message.reply('❌ Error happened while processing your request', parse_mode='html',
reply_markup=get_exception_keyboard(exception_id))
elif isinstance(message, types.CallbackQuery):
await message.message.reply('❌ Error happened while processing your request',
parse_mode='html',
reply_markup=get_exception_keyboard(exception_id))
return
else:
raise e
return wrapper
return decorator