Update bot/telegram/media.py

bot/telegram/media.py

@@ -1,22 +1,64 @@
 # PATH: bot/telegram/media.py
 import os
 import time
-from typing import Callable, Optional, Tuple
+import asyncio
+import inspect
+from typing import Callable, Any, Optional, Tuple
+
 from hydrogram import Client
 from hydrogram.types import Message
+
 from bot.temp.files import make_temp_file_path
 
-ProgressCb = Callable[[int, int], None]
+
+def _wrap_progress(cb: Callable[..., Any], extra_args: Tuple[Any, ...] = ()):
+    """
+    Hydrogram/Pyrogram progress callback is usually sync:
+      progress(current, total, *args)
+
+    If user passes async cb, we schedule it.
+    """
+    if not cb:
+        return None, ()
+
+    async_cb = inspect.iscoroutinefunction(cb)
+
+    def _progress(current: int, total: int, *args):
+        try:
+            cur = int(current or 0)
+            tot = int(total or 0)
+        except Exception:
+            cur, tot = 0, 0
+
+        if async_cb:
+            try:
+                asyncio.create_task(cb(cur, tot, *extra_args))
+            except Exception:
+                pass
+        else:
+            try:
+                cb(cur, tot, *extra_args)
+            except Exception:
+                pass
+
+    return _progress, ()
+
 
 async def download_to_temp(
     app: Client,
     m: Message,
-    progress: Optional[ProgressCb] = None,
-    progress_args: tuple = (),
-) -> Tuple[str, int, str, float]:
+    progress_cb: Optional[Callable[..., Any]] = None,
+    progress_args: Tuple[Any, ...] = (),
+    return_meta: bool = False,
+):
     """
-    Returns: (file_path, file_size, file_name, elapsed_sec)
-    progress(done, total) will be called if supported by Hydrogram.
+    Download Telegram media (video/document) to a temp path.
+
+    Returns (default):
+      (file_path, file_size, file_name)
+
+    If return_meta=True:
+      (file_path, file_size, file_name, download_seconds)
     """
     media = m.video or m.document
     if not media:
@@ -28,19 +70,23 @@ async def download_to_temp(
     out_path = make_temp_file_path(file_name)
     os.makedirs(os.path.dirname(out_path), exist_ok=True)
 
-    t0 = time.time()
+    prog, _ = _wrap_progress(progress_cb, progress_args)
 
-    # Hydrogram (Pyrogram-like) usually supports progress/progress_args.
-    # If not supported, fall back.
+    t0 = time.monotonic()
     try:
+        # Hydrogram usually matches Pyrogram's signature:
+        # download_media(message, file_name=..., progress=..., progress_args=...)
         await app.download_media(
             message=m,
             file_name=out_path,
-            progress=progress,
+            progress=prog,
             progress_args=progress_args,
         )
     except TypeError:
+        # Fallback if this Hydrogram build doesn't support progress args
         await app.download_media(message=m, file_name=out_path)
+    t1 = time.monotonic()
 
-    elapsed = max(0.001, time.time() - t0)
-    return out_path, file_size, file_name, elapsed
+    if return_meta:
+        return out_path, file_size, file_name, (t1 - t0)
+    return out_path, file_size, file_name