Decorator error handling

bot/callbacks/set_model.py

@@ -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...')
+    await call.message.edit_reply_markup(reply_markup=None)
-    try:
     await models.set_model(db[DBTables.config]['models'][n])
 
-    except ClientConnectorError:
+    await call.message.answer('✅ Model set for all users!')
-        await call.answer('❌ Error! Maybe, StableDiffusion API endpoint is incorrect '
-                          'or turned off', show_alert=True)
     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):

bot/handlers/admin/tools.py

@@ -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.errorable_command import wrap_exception
-from bot.utils.trace_exception import PrettyException
 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

bot/handlers/image_info/image_info.py

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

bot/handlers/txt2img/set_model.py

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

bot/handlers/txt2img/set_settings.py

@@ -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.errorable_command import wrap_exception
-from bot.utils.trace_exception import PrettyException
 
 
+@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')
 

bot/handlers/txt2img/txt2img.py

@@ -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 bot.utils.errorable_command import wrap_exception
-from aiohttp import ClientConnectorError
 
 
+@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

bot/utils/errorable_command.py

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