Better task management

wbb/__init__.py

@@ -26,11 +26,13 @@ print("[INFO]: INITIALIZING")
 import asyncio
 import logging
 import time
+from inspect import getfullargspec
 from os import path
 
 from aiohttp import ClientSession
 from motor.motor_asyncio import AsyncIOMotorClient as MongoClient
 from pyrogram import Client
+from pyrogram.types import Message
 from pyromod import listen
 from Python_ARQ import ARQ
 
@@ -171,3 +173,11 @@ if USERBOT_ID not in SUDOERS:
     SUDOERS.append(USERBOT_ID)
 app.stop()
 app2.stop()
+
+
+async def eor(msg: Message, **kwargs):
+    func = msg.edit_text if msg.from_user.is_self else msg.reply
+    spec = getfullargspec(func.__wrapped__).args
+    return await func(
+        **{k: v for k, v in kwargs.items() if k in spec}
+    )

wbb/core/tasks.py

@@ -0,0 +1,121 @@
+from asyncio import Lock, create_task
+from time import time
+
+from pyrogram import filters
+from pyrogram.types import Message
+
+from wbb import SUDOERS, USERBOT_PREFIX, app2, eor
+from wbb.core.sections import bold, section, w
+
+tasks = {}
+TASKS_LOCK = Lock()
+arrow = lambda x: x.text + "\n`→`"
+
+
+def all_tasks():
+    return tasks
+
+
+async def add_task(
+    taskFunc,
+    task_name,
+    *args,
+    **kwargs,
+):
+
+    async with TASKS_LOCK:
+        global tasks
+
+        task_id = (list(tasks.keys())[-1] + 1) if tasks else 0
+
+        task = create_task(
+            taskFunc(*args, **kwargs),
+            name=task_name,
+        )
+        tasks[task_id] = task, int(time())
+    return task, task_id
+
+
+async def rm_task(task_id=None):
+    global tasks
+
+    async with TASKS_LOCK:
+        for key, value in list(tasks.items()):
+            if value[0].done() or value[0].cancelled():
+                del tasks[key]
+
+        if task_id:
+            if task_id in tasks:
+                if not tasks[task_id][0].done():
+                    tasks[task_id][0].cancel()
+                del tasks[task_id]
+
+
+@app2.on_message(
+    filters.user(SUDOERS)
+    & ~filters.forwarded
+    & ~filters.via_bot
+    & filters.command("cancelTask", prefixes=USERBOT_PREFIX)
+)
+async def task_cancel(_, message: Message):
+    m = message
+    r = m.reply_to_message
+
+    if len(m.text.split()) == 2:
+        mid = int(m.text.split(None, 1)[1])
+    else:
+        mid = r.message_id if r else None
+
+    tasks = all_tasks()
+
+    if not mid or not tasks:
+        return await m.delete()
+
+    if mid not in tasks:
+        return await m.delete()
+
+    await rm_task(mid)
+    await eor(message, text=f"{arrow(m)} Task cancelled")
+
+
+@app2.on_message(
+    filters.user(SUDOERS)
+    & ~filters.forwarded
+    & ~filters.via_bot
+    & filters.command("lsTasks", prefixes=USERBOT_PREFIX)
+)
+async def task_list(_, message: Message):
+    await rm_task()  # Clean completed tasks
+
+    tasks = all_tasks()
+
+    if not tasks:
+        return await eor(
+            message,
+            text=f"{arrow(message)} No tasks pending",
+        )
+
+    text = bold("Tasks") + "\n"
+
+    for i, task in enumerate(list(tasks.items())):
+        indent = w * 4
+
+        t, started = task[1]
+        elapsed = round(time() - started)
+        info = t._repr_info()
+
+        id = task[0]
+
+        text += section(
+            f"{indent}Task {i}",
+            body={
+                "Name": t.get_name(),
+                "Task ID": id,
+                "Status": info[0].capitalize(),
+                "Origin": info[2].split("/")[-1].replace(">", ""),
+                "Running since": f"{elapsed}s",
+            },
+            indent=6,
+        )
+
+    await eor(message, text=text)

wbb/modules/chatbot.py

@@ -27,9 +27,8 @@ from pyrogram import filters
 from pyrogram.types import Message
 
 from wbb import (BOT_ID, SUDOERS, USERBOT_ID, USERBOT_PREFIX,
-                 USERBOT_USERNAME, app, app2, arq)
+                 USERBOT_USERNAME, app, app2, arq, eor)
 from wbb.core.decorators.errors import capture_err
-from wbb.modules.userbot import eor
 from wbb.utils.filter_groups import chatbot_group
 
 __MODULE__ = "ChatBot"

wbb/modules/devs.py

@@ -18,7 +18,7 @@ from pyrogram import filters
 from pyrogram.types import (InlineKeyboardButton,
                             InlineKeyboardMarkup, Message)
 
-from wbb import SUDOERS, app
+from wbb import SUDOERS, app, eor
 
 # Module help for this is in sudoers.py
 
@@ -31,12 +31,6 @@ async def aexec(code, client, message):
     return await locals()["__aexec"](client, message)
 
 
-async def eor(msg: Message, **kwargs):
-    func = msg.edit_text if msg.from_user.is_self else msg.reply
-    spec = getfullargspec(func.__wrapped__).args
-    await func(**{k: v for k, v in kwargs.items() if k in spec})
-
-
 @app.on_message(
     filters.user(SUDOERS)
     & ~filters.forwarded

wbb/modules/download_upload.py

@@ -1,13 +1,39 @@
+"""
+MIT License
+
+Copyright (c) 2021 TheHamkerCat
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+
 from os import remove
 from os.path import isfile
 from time import ctime, time
+from traceback import format_exc
 
 from pyrogram import filters
 from pyrogram.types import Message
 
-from wbb import SUDOERS, USERBOT_PREFIX, app2
+from wbb import SUDOERS, USERBOT_PREFIX, app2, eor
 from wbb.core.sections import section
-from wbb.modules.userbot import add_task, eor, rm_task
+from wbb.core.tasks import add_task, rm_task
+from wbb.modules.userbot import eor
 from wbb.utils.downloader import download
 
 
@@ -18,11 +44,9 @@ from wbb.utils.downloader import download
 async def download_func(_, message: Message):
     reply = message.reply_to_message
     start = time()
-    task_id = int(start)
 
     body = {
         "Started": ctime(start),
-        "Task ID": task_id,
     }
     m = await eor(
         message,
@@ -30,15 +54,24 @@ async def download_func(_, message: Message):
     )
 
     if reply:
-        task = await add_task(
+        task, task_id = await add_task(
             reply.download,
-            task_id=task_id,
             task_name="Downloader",
         )
+
+        body["Task ID"] = task_id
+        await eor(
+            m,
+            text=section("Downloading", body),
+        )
+
         await task
         await rm_task(task_id)
 
+        file = task.result()
+
         elapsed = int(time() - start)
+        body["File"] = file.split("/")[-1]
         body["Took"] = f"{elapsed}s"
 
         return await eor(m, text=section("Downloaded", body))
@@ -50,14 +83,25 @@ async def download_func(_, message: Message):
     url = text.split(None, 1)[1]
 
     try:
-        await download(
-            url,
-            task_id=task_id,
+        task, task_id = await download(url)
+
+        body["Task ID"] = task_id
+        await eor(
+            m,
+            text=section("Downloading", body),
         )
+
+        await task
+        file = task.result()
+        await rm_task(task_id)
+
     except Exception as e:
+        e = format_exc()
+        e = e.splitlines()[-1]
         return await eor(m, text=f"**Error:** `{str(e)}`")
 
     elapsed = int(time() - start)
+    body["File"] = file.split("/")[-1]
     body["Took"] = f"{elapsed}s"
 
     await eor(m, text=section("Downloaded", body))
@@ -74,23 +118,26 @@ async def upload_func(_, message: Message):
     url_or_path = message.text.split(None, 1)[1]
 
     start = time()
-    task_id = int(start)
 
     body = {
         "Started": ctime(start),
-        "Task ID": task_id,
     }
 
     m = await eor(message, text=section("Uploading", body))
 
     async def upload_file(path: str):
-        task = await add_task(
+        task, task_id = await add_task(
             message.reply_document,
-            task_id,
             "Uploader",
             path,
         )
 
+        body["Task ID"] = task_id
+        await eor(
+            m,
+            text=section("Uploading", body),
+        )
+
         await task
         await rm_task(task_id)
 
@@ -103,11 +150,20 @@ async def upload_func(_, message: Message):
         if isfile(url_or_path):
             return await upload_file(url_or_path)
 
-        path = await download(
-            url_or_path,
-            task_id=task_id,
+        task, task_id = await download(url_or_path)
+
+        body["Task ID"] = task_id
+        await eor(
+            m,
+            text=section("Downloading", body),
         )
+        await task
+        file = task.result()
+        await rm_task(task_id)
 
-        return await upload_file(path)
+        await upload_file(file)
+        remove(file)
     except Exception as e:
+        e = format_exc()
+        e = e.splitlines()[-1]
         return await eor(m, text=f"**Error:** `{str(e)}`")

wbb/modules/img_pdf.py

@@ -84,6 +84,7 @@ async def convert(
 
 
 @app.on_message(filters.command("pdf"))
+@capture_err
 async def img_to_pdf(_, message: Message):
     reply = message.reply_to_message
     if not reply:

wbb/modules/pmpermit.py

@@ -25,9 +25,9 @@ SOFTWARE.
 from pyrogram import filters
 from pyrogram.raw.functions.messages import DeleteHistory
 
-from wbb import BOT_ID, SUDOERS, USERBOT_ID, USERBOT_PREFIX, app, app2
+from wbb import (BOT_ID, SUDOERS, USERBOT_ID, USERBOT_PREFIX, app,
+                 app2, eor)
 from wbb.core.decorators.errors import capture_err
-from wbb.modules.userbot import eor
 from wbb.utils.dbfunctions import (approve_pmpermit,
                                    disapprove_pmpermit,
                                    is_pmpermit_approved)

wbb/modules/reverse.py

@@ -100,7 +100,10 @@ async def reverse_image_search(_, message: Message):
         text = div.find("a").text
         text = f"**Result**: [{text}]({location})"
     except Exception:
-        return await m.edit(f"**Result**: [Link]({location})")
+        return await m.edit(
+            f"**Result**: [Link]({location})",
+            disable_web_page_preview=True,
+        )
 
     # Pass if no images detected
     try:
@@ -145,4 +148,7 @@ async def reverse_image_search(_, message: Message):
     except Exception:
         pass
 
-    await m.edit(text)
+    await m.edit(
+        f"**Result**: [{text}]({location})",
+        disable_web_page_preview=True,
+    )

wbb/modules/sudo.py

@@ -24,9 +24,8 @@ SOFTWARE.
 from pyrogram import filters
 from pyrogram.types import Message
 
-from wbb import BOT_ID, SUDOERS, USERBOT_PREFIX, app2
+from wbb import BOT_ID, SUDOERS, USERBOT_PREFIX, app2, eor
 from wbb.core.decorators.errors import capture_err
-from wbb.modules.userbot import eor
 from wbb.utils.dbfunctions import add_sudo, get_sudoers, remove_sudo
 from wbb.utils.functions import restart
 

wbb/modules/userbot.py

@@ -10,18 +10,17 @@ import re
 import subprocess
 import sys
 import traceback
-from asyncio import Lock, create_task
 from html import escape
-from inspect import getfullargspec
 from io import StringIO
 
 from pyrogram import filters
 from pyrogram.errors import MessageNotModified
 from pyrogram.types import Message, ReplyKeyboardMarkup
 
-from wbb import app  # don't remove
-from wbb import SUDOERS, USERBOT_PREFIX, app2, arq
+from wbb import (SUDOERS, USERBOT_PREFIX, app, app2,  # don't remove
+                 arq, eor)
 from wbb.core.sections import section
+from wbb.core.tasks import add_task, rm_task
 
 # Eval and Sh module from nana-remix
 
@@ -29,9 +28,7 @@ m = None
 p = print
 r = None
 arq = arq
-arrow = lambda x: x.text + "\n`→`"
-TASKS_LOCK = Lock()
-tasks = {}
+arrow = lambda x: (x.text if isinstance(x, Message) else "") + "\n`→`"
 
 
 async def aexec(code, client, message):
@@ -42,83 +39,24 @@ async def aexec(code, client, message):
     return await locals()["__aexec"](client, message)
 
 
-async def eor(msg: Message, **kwargs):
-    func = msg.edit_text if msg.from_user.is_self else msg.reply
-    spec = getfullargspec(func.__wrapped__).args
-    return await func(
-        **{k: v for k, v in kwargs.items() if k in spec}
-    )
-
-
-async def add_task(taskFunc, task_id, task_name, *args, **kwargs):
-    global tasks
-    task = create_task(
-        taskFunc(*args, **kwargs),
-        name=task_name,
-    )
-    tasks[task_id] = task
-    return task
-
-
-async def rm_task(task_id=None):
-    global tasks
-    async with TASKS_LOCK:
-        for key, value in list(tasks.items()):
-            if value.done() or value.cancelled():
-                del tasks[key]
-
-        if task_id:
-            if task_id in tasks:
-                if not tasks[task_id].done():
-                    tasks[task_id].cancel()
-                del tasks[task_id]
-
-
-@app2.on_message(
-    filters.user(SUDOERS)
-    & ~filters.forwarded
-    & ~filters.via_bot
-    & filters.command("cancelTask", prefixes=USERBOT_PREFIX)
-)
-async def task_cancel(_, message: Message):
-    m = message
-    r = m.reply_to_message
+async def iter_edit(message: Message, text: str):
+    async for m in app2.iter_history(message.chat.id):
 
-    if len(m.text.split()) == 2:
-        mid = int(m.text.split(None, 1)[1])
-    else:
-        mid = r.message_id if r else None
+        # If no replies found, reply
+        if m.message_id == message.message_id:
+            return 0
 
-    if not mid or not tasks:
-        return await m.delete()
+        if not m.from_user or not m.text or not m.reply_to_message:
+            continue
 
-    if mid not in tasks:
-        return await m.delete()
+        if m.reply_to_message.message_id == message.message_id:
+            if m.from_user.id == message.from_user.id:
 
-    await rm_task(mid)
-    await eor(message, text=f"{arrow(m)} Task cancelled")
-
-
-@app2.on_message(
-    filters.user(SUDOERS)
-    & ~filters.forwarded
-    & ~filters.via_bot
-    & filters.command("lsTasks", prefixes=USERBOT_PREFIX)
-)
-async def task_list(_, message: Message):
-    await rm_task()
-    if not tasks:
-        return await eor(
-            message,
-            text=f"{arrow(message)} No tasks pending",
-        )
-
-    body = {
-        str(key): value.get_name()
-        for key, value in list(tasks.items())
-    }
-
-    await eor(message, text=section("Pending Tasks", body))
+                if "→" in m.text:
+                    try:
+                        return await m.edit(text)
+                    except MessageNotModified:
+                        return
 
 
 @app2.on_message(
@@ -146,23 +84,30 @@ async def executor(client, message: Message):
         if r.reply_markup:
             if isinstance(r.reply_markup, ReplyKeyboardMarkup):
                 return await eor(m, text="INSECURE!")
+    status = None
     old_stderr = sys.stderr
     old_stdout = sys.stdout
     redirected_output = sys.stdout = StringIO()
     redirected_error = sys.stderr = StringIO()
     stdout, stderr, exc = None, None, None
     try:
-        task = await add_task(
+        task, task_id = await add_task(
             aexec,
-            m.message_id,
             "Eval",
             cmd,
             client,
-            message,
+            m,
         )
+
+        text = f"{arrow('')} Pending Task `{task_id}`"
+        if not message.edit_date:
+            status = await m.reply(text, quote=True)
+
         await task
     except Exception as e:
-        exc = str(e)
+        e = traceback.format_exc()
+        print(e)
+        exc = e.splitlines()[-1]
 
     await rm_task()
 
@@ -182,7 +127,9 @@ async def executor(client, message: Message):
         globals()["lstdout"] = stdout
     else:
         evaluation = "Success"
+
     final_output = f"**→**\n`{escape(evaluation.strip())}`"
+
     if len(final_output) > 4096:
         filename = "output.txt"
         with open(filename, "w+", encoding="utf8") as out_file:
@@ -193,33 +140,16 @@ async def executor(client, message: Message):
             quote=False,
         )
         os.remove(filename)
-    else:
-        mid = message.message_id
-
-        # Edit the output if input is edited
-        if message.edit_date:
-            async for m in app2.iter_history(message.chat.id):
-
-                # If no replies found, reply
-                if m.message_id == mid:
-                    break
-
-                if (
-                    not m.from_user
-                    or not m.text
-                    or not m.reply_to_message
-                ):
-                    continue
-
-                if m.reply_to_message.message_id == mid:
-                    if m.from_user.id == message.from_user.id:
-
-                        if "→" in m.text:
-                            try:
-                                return await m.edit(final_output)
-                            except MessageNotModified:
-                                return
-        await message.reply(final_output, quote=True)
+        return await status.delete()
+
+    # Edit the output if input is edited
+    if message.edit_date:
+        status_ = await iter_edit(message, final_output)
+        if status_ == 0:
+            pass
+        else:
+            return
+    await eor(status, text=final_output, quote=True)
 
 
 @app2.on_message(

wbb/utils/downloader.py

@@ -5,7 +5,7 @@ from time import time
 import aiofiles
 
 from wbb import aiohttpsession as session
-from wbb.modules.userbot import add_task, rm_task
+from wbb.core.tasks import add_task, rm_task
 
 
 def ensure_status(status_code: int):
@@ -44,18 +44,20 @@ async def download(
     :url: url where the file is located
     :file_path: path/to/file
     :chunk_size: size of a single chunk
+
+    Returns:
+            (asyncio.Task, task_id), With which you can await
+            the task, track task progress or cancel it.
     """
     # Create a task and add it to main tasks dict
     # So we can cancel it using .cancelTask
 
-    task = await add_task(
+    task, task_id = await add_task(
         download_url,
-        task_id,
         "Downloader",
         url=url,
         file_path=file_path,
         chunk_size=chunk_size,
     )
-    await task
-    await rm_task(task_id)
-    return task.result()
+
+    return task, task_id