feat: stop running message execution on cancel, add stop command …

…to only stop execution without blacklisting the message

tgpy/_core/eval_message.py

@@ -1,3 +1,6 @@
+import asyncio
+from asyncio import Task
+
 from telethon.errors import MessageIdInvalidError
 from telethon.tl.custom import Message
 
@@ -6,29 +9,34 @@
 from tgpy._core.utils import convert_result, format_traceback
 from tgpy.api import constants, tgpy_eval
 
+running_messages: dict[tuple[int, int], Task] = {}
 
-async def eval_message(code: str, message: Message) -> Message | None:
-    await message_design.edit_message(message, code, 'Running...')
 
+async def eval_message(code: str, message: Message) -> Message | None:
+    task = asyncio.create_task(tgpy_eval(code, message, filename=None))
+    running_messages[(message.chat_id, message.id)] = task
     # noinspection PyBroadException
     try:
-        eval_result = await tgpy_eval(code, message, filename=None)
+        eval_result = await task
+    except asyncio.CancelledError:
+        # message cancelled, do nothing
+        # return no message as it wasn't edited
+        return None
     except Exception:
         result = None
         output = ''
-        exc, full_exc = format_traceback()
+        exc, constants['exc'] = format_traceback()
     else:
         if app.ctx.is_manual_output:
             return
         result = convert_result(eval_result.result)
         output = eval_result.output
         exc = ''
-        full_exc = None
-
-    constants['exc'] = full_exc
+        constants['exc'] = None
+    finally:
+        running_messages.pop((message.chat_id, message.id))
 
     try:
-        # noinspection PyProtectedMember
         return await message_design.edit_message(
             message,
             code,
@@ -40,4 +48,4 @@ async def eval_message(code: str, message: Message) -> Message | None:
         return None
 
 
-__all__ = ['eval_message']
+__all__ = ['eval_message', 'running_messages']

tgpy/std/prevent_eval.py

@@ -10,45 +10,51 @@
 from telethon.tl.custom import Message
 
 import tgpy.api
-from tgpy import Context
+from tgpy import Context, reactions_fix
+from tgpy._core.eval_message import running_messages
 
 client: TelegramClient
 ctx: Context
 
 MODULE_NAME = 'prevent_eval'
 IGNORED_MESSAGES_KEY = f'{MODULE_NAME}.ignored_messages'
 CANCEL_RGX = re.compile(r'(?i)^(cancel|сфтсуд)$')
+INTERRUPT_RGX = re.compile(r'(?i)^(stop|ыещз)$')
 
 
-async def handle_cancel(message: Message):
-    prev: Message = await message.get_reply_message()
-    if prev:
-        parsed = tgpy.api.parse_tgpy_message(prev)
+async def cancel_message(message: Message, permanent: bool = True) -> None:
+    parsed = tgpy.api.parse_tgpy_message(message)
+
+    if task := running_messages.get((message.chat_id, message.id)):
+        task.cancel()
+    message = await message.edit(parsed.code)
+
+    if permanent:
+        ignored_messages = tgpy.api.config.get(IGNORED_MESSAGES_KEY, [])
+        ignored_messages.append([message.chat_id, message.id])
+        tgpy.api.config.save()
     else:
+        reactions_fix.update_hash(message)
+
+
+async def handle_cancel(message: Message, permanent: bool = True):
+    target: Message = await message.get_reply_message()
+    if not target:
         async for msg in client.iter_messages(
             message.chat_id, max_id=message.id, limit=10
         ):
             if not msg.out:
                 continue
             parsed = tgpy.api.parse_tgpy_message(msg)
             if parsed.is_tgpy_message:
-                prev = msg
+                target = msg
                 break
-        else:
-            return
-
-    ignored_messages = tgpy.api.config.get(IGNORED_MESSAGES_KEY, [])
-    ignored_messages.append([prev.chat_id, prev.id])
-    tgpy.api.config.save()
 
-    # noinspection PyBroadException
-    try:
-        edit_result = await prev.edit(parsed.code)
-    except Exception:
+    if not target:
         return
 
-    if edit_result is not None:
-        await message.delete()
+    await cancel_message(target, permanent)
+    await message.delete()
 
 
 async def handle_comment(message: Message):
@@ -73,11 +79,12 @@ async def exec_hook(message: Message, is_edit: bool):
 
     is_comment = message.raw_text.startswith('//') and message.raw_text[2:].strip()
     is_cancel = CANCEL_RGX.fullmatch(message.raw_text) is not None
-    if not is_comment and not is_cancel:
+    is_interrupt = INTERRUPT_RGX.fullmatch(message.raw_text) is not None
+    if not is_comment and not is_cancel and not is_interrupt:
         return True
 
-    if is_cancel:
-        await handle_cancel(message)
+    if is_cancel or is_interrupt:
+        await handle_cancel(message, permanent=is_cancel)
     elif is_comment:
         await handle_comment(message)