Add exec tasks

sukuinote/plugins/pyexec.py

@@ -9,6 +9,8 @@ from io import StringIO, BytesIO
 from pyrogram import Client, filters
 from .. import config, help_dict, log_errors, slave, apps, session, public_log_errors
 
+exec_tasks = dict()
+
 PYEXEC_REGEX = '^(?:' + '|'.join(map(re.escape, config['config']['prefixes'])) + r')exec\s+([\s\S]+)$'
 @Client.on_message(~filters.forwarded & ~filters.sticker & ~filters.via_bot & ~filters.edited & filters.me & filters.regex(PYEXEC_REGEX))
 @log_errors
@@ -39,8 +41,15 @@ async def pyexec(client, message):
         if ex.msg != "'return' with value in async generator":
             raise
         exx = _gf(obody)
-    reply = await message.reply_text('Executing...')
-    async_obj = exx(client, client, message, message, reply, message.reply_to_message, message.reply_to_message, UniqueExecReturnIdentifier)
+    reply = await message.reply_text(f'Executing <code>{hash(UniqueExecReturnIdentifier)}</code>...')
+    oasync_obj = exx(client, client, message, message, reply, message.reply_to_message, message.reply_to_message, UniqueExecReturnIdentifier)
+    if inspect.isasyncgen(oasync_obj):
+        async def async_obj():
+            return [i async for i in oasync_obj]
+    else:
+        async def async_obj():
+            to_return = [await oasync_obj]
+            return [] if to_return == [UniqueExecReturnIdentifier] else to_return
     stdout = sys.stdout
     stderr = sys.stderr
     wrapped_stdout = StringIO()
@@ -48,15 +57,19 @@ async def pyexec(client, message):
     try:
         sys.stdout = wrapped_stdout
         sys.stderr = wrapped_stderr
-        if inspect.isasyncgen(async_obj):
-            returned = [i async for i in async_obj]
-        else:
-            returned = [await async_obj]
-            if returned == [UniqueExecReturnIdentifier]:
-                returned = []
+        task = asyncio.create_task(async_obj())
+        exec_tasks[hash(UniqueExecReturnIdentifier)] = task
+        returned = await async_obj
+    except asyncio.CancelledError:
+        sys.stdout = stdout
+        sys.stderr = stderr
+        exec_tasks.pop(hash(UniqueExecReturnIdentifier), None)
+        await reply.edit_text('Cancelled')
+        return
     finally:
         sys.stdout = stdout
         sys.stderr = stderr
+        exec_tasks.pop(hash(UniqueExecReturnIdentifier), None)
     wrapped_stderr.seek(0)
     wrapped_stdout.seek(0)
     output = ''
@@ -83,4 +96,36 @@ async def pyexec(client, message):
     else:
         await reply.edit_text(output)
 
-help_dict['exec'] = ('Exec', '{prefix}exec <i>&lt;python code&gt;</i> - Executes python code')
+@Client.on_message(~filters.forwarded & ~filters.sticker & ~filters.via_bot & ~filters.edited & filters.me & filters.command(['listexecs', 'listexec', 'lexec'], prefixes=config['config']['prefixes']))
+@log_errors
+@public_log_errors
+async def listexec(client, message):
+    text = '\n'.join(map(str, exec_tasks))
+    if len(text) > 4096:
+        f = BytesIO(text.encode('utf-8'))
+        f.name = 'exectasks.txt'
+        await message.reply_document(f)
+    else:
+        text = '\n'.join(map(lambda i: f'<code>{i}</code>', exec_tasks))
+        await message.reply_text(text)
+
+@Client.on_message(~filters.forwarded & ~filters.sticker & ~filters.via_bot & ~filters.edited & filters.me & filters.command(['cancelexec', 'cexec'], prefixes=config['config']['prefixes']))
+@log_errors
+@public_log_errors
+async def cancelexec(client, message):
+    try:
+        task = exec_tasks.get(int(message[1]))
+    except IndexError:
+        return
+    if not task:
+        await message.reply_text('Task does not exist')
+        return
+    task.cancel()
+
+help_dict['exec'] = ('Exec', '''{prefix}exec <i>&lt;python code&gt;</i> - Executes python code
+
+{prefix}listexecs - List exec tasks
+Aliases: {prefix}listexec, {prefix}lexec
+
+{prefix}cancelexec <i>&lt;task id&gt;</i> - Cancel exec task
+Aliases: {prefix}cexec''')