Upload small files (100MB) using ARQ /storage

wbb/modules/arq_storage.py

@@ -0,0 +1,68 @@
+from asyncio import Lock
+from os import remove
+
+from pyrogram import filters
+from pyrogram.types import Message
+
+from wbb import app, arq
+from wbb.core.decorators.errors import capture_err
+from wbb.utils.functions import get_file_id_from_message
+
+__MODULE__ = "Storage"
+__HELP__ = """
+Upload unlimited files smaller than 100MB
+And get a download link
+
+**Usage:**
+    /upload [url|Reply to a file]
+"""
+
+UPLOAD_LOCK = Lock()
+
+
+async def upload(m: Message, file: str = None, url: str = None):
+    err = "Something went wrong"
+    try:
+        resp = await arq.upload(file=file, url=url)
+    except Exception:
+        await m.edit(err)
+        if file:
+            remove(file)
+        return
+    if not resp:
+        return await m.edit(err)
+    if not resp.ok:
+        return await m.edit(err)
+    await m.edit(
+        f"**Download Link:** {resp.result}",
+        disable_web_page_preview=True,
+    )
+
+
+@app.on_message(filters.command("upload"))
+async def arq_upload(_, message):
+    if message.reply_to_message:
+        if UPLOAD_LOCK.locked():
+            return await message.reply(
+                "One upload is already in progress, please try again later"
+            )
+        async with UPLOAD_LOCK:
+            r = message.reply_to_message
+
+            file_id = get_file_id_from_message(r, 100000000, None)
+            if not file_id:
+                return await message.reply("Media not supported.")
+
+            m = await message.reply("Downloading...")
+            file = await app.download_media(file_id)
+
+            await m.edit("Uploading...")
+            return await upload(m, file=file)
+
+    if len(message.command) != 2:
+        return await message.reply("Not enough arguments")
+
+    url = message.text.split(None, 1)[1]
+
+    m = await message.reply("Uploading...")
+    await upload(m, url=url)

wbb/modules/reverse.py

@@ -32,7 +32,7 @@ async def reverse_image_search(_, message):
             "Reply to an image/document/sticker/animation to reverse search it."
         )
     m = await message.reply_text("Downloading")
-    file_id = await get_file_id_from_message(reply)
+    file_id = get_file_id_from_message(reply)
     if not file_id:
         return await m.edit("Can't reverse that")
     image = await app.download_media(

wbb/utils/functions.py

@@ -273,54 +273,19 @@ async def extract_user(message):
     return (await extract_user_and_reason(message))[0]
 
 
-async def test_ARQ(message):
-    results = ""
-    funcs = {
-        "image": arq.image("something"),
-        "luna": arq.luna("hello"),
-        "lyrics": arq.lyrics("attention"),
-        "nlp": arq.nlp("bitcoin"),
-        "nsfw_scan": arq.nsfw_scan(
-            url="https://www.pixsy.com/wp-content/uploads/2021/04/ben-sweet-2LowviVHZ-E-unsplash-1.jpeg"
-        ),
-        "pornhub": arq.pornhub("something"),
-        "proxy": arq.proxy(),
-        "pypi": arq.pypi("python-arq"),
-        "reddit": arq.reddit("badcode"),
-        "quotly": arq.quotly(message),
-        "saavn": arq.saavn("attention"),
-        "stats": arq.stats(),
-        "tmdb": arq.tmdb("flash"),
-        "torrent": arq.torrent("porn"),
-        "translate": arq.translate("hello"),
-        "urbandict": arq.urbandict("wtf"),
-        "wall": arq.wall("anime"),
-        "wiki": arq.wiki("cat"),
-        "youtube": arq.youtube("never gonna give you up"),
-    }
-
-    for key, value in funcs.items():
-        try:
-            t1 = time()
-            result = await value
-            t2 = time()
-            if result.ok:
-                results += f"**{key.capitalize()}:** `{t2-t1}`\n"
-            else:
-                results += f"**{key.capitalize()}:** `Failed`\n"
-        except Exception:
-            results += f"**{key.capitalize()}:** `Failed`\n"
-    return results
-
-
-async def get_file_id_from_message(message):
+def get_file_id_from_message(
+    message,
+    max_file_size=3145728,
+    mime_types=["image/png", "image/jpeg"],
+):
     file_id = None
     if message.document:
-        if int(message.document.file_size) > 3145728:
+        if int(message.document.file_size) > max_file_size:
             return
         mime_type = message.document.mime_type
-        if mime_type != "image/png" and mime_type != "image/jpeg":
-            return
+        if mime_types:
+            if mime_type not in mime_types:
+                return
         file_id = message.document.file_id
 
     if message.sticker: