diff --git a/sukuinote/plugins/pyexec.py b/sukuinote/plugins/pyexec.py
index d2375dd..aa0ee36 100644
--- a/sukuinote/plugins/pyexec.py
+++ b/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 {hash(UniqueExecReturnIdentifier)}
...')
+ 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 <python code> - 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'{i}
', 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 <python code> - Executes python code
+
+{prefix}listexecs - List exec tasks
+Aliases: {prefix}listexec, {prefix}lexec
+
+{prefix}cancelexec <task id> - Cancel exec task
+Aliases: {prefix}cexec''')