From f538e143ba6f1c8e5c6db34f38a88ef5bacecec7 Mon Sep 17 00:00:00 2001
From: BennyThink <benny.think@gmail.com>
Date: Wed, 5 May 2021 14:56:47 +0800
Subject: [PATCH] add 2GiB warning, add total size to download notice

---
 ytdl.py | 51 +++++++++++++++++++++++++++++++--------------------
 1 file changed, 31 insertions(+), 20 deletions(-)

diff --git a/ytdl.py b/ytdl.py
index cf492fd..fc16120 100644
--- a/ytdl.py
+++ b/ytdl.py
@@ -26,7 +26,6 @@ from hachoir.metadata.video import MkvMetadata
 from telethon import TelegramClient, events
 from telethon.tl.types import DocumentAttributeFilename, DocumentAttributeVideo
 from telethon.utils import get_input_media
-
 from tgbot_ping import get_runtime
 
 from FastTelethon.FastTelethon import upload_file
@@ -63,34 +62,29 @@ def get_metadata(video_path):
         return dict(duration=0, w=0, h=0), 'application/octet-stream'
 
 
-async def upload_callback(current, total, chat_id, message):
-    key = f"{chat_id}-{message.id}"
-    # if the key exists, we shouldn't send edit message
-    if not r.exists(key):
-        r.set(key, "ok", ex=EXPIRE)
-        msg = f'Uploading {round(current / total * 100, 2)}%: {current}/{total}'
-        await bot.edit_message(chat_id, message, msg)
-
-
-async def sync_edit_message(chat_id, message, msg):
-    # try to avoid flood
-    key = f"{chat_id}-{message.id}"
-    if not r.exists(key):
-        r.set(key, "ok", ex=EXPIRE)
-        await bot.edit_message(chat_id, message, msg)
-
-
 def go(chat_id, message, msg):
     asyncio.run(sync_edit_message(chat_id, message, msg))
 
 
+def sizeof_fmt(num: int, suffix='B'):
+    for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
+        if abs(num) < 1024.0:
+            return "%3.1f%s%s" % (num, unit, suffix)
+        num /= 1024.0
+    return "%.1f%s%s" % (num, 'Yi', suffix)
+
+
 def progress_hook(d: dict, chat_id, message):
     if d['status'] == 'downloading':
         downloaded = d.get("downloaded_bytes", 0)
-        total = d.get("total_bytes") or d.get("total_bytes_estimate", "N/A")
+        total = d.get("total_bytes") or d.get("total_bytes_estimate", 0)
+        filesize = sizeof_fmt(total)
+        if total > 2 * 1024 * 1024 * 1024:
+            raise Exception("\n\nYour video is too large. %s will exceed Telegram's max limit 2GiB" % filesize)
+
         percent = d.get("_percent_str", "N/A")
         speed = d.get("_speed_str", "N/A")
-        msg = f'Downloading {percent}: {downloaded}/{total} @ {speed}'
+        msg = f'[{filesize}]: Downloading {percent} - {downloaded}/{total} @ {speed}'
         threading.Thread(target=go, args=(chat_id, message, msg)).start()
 
 
@@ -127,6 +121,23 @@ def ytdl_download(url, tempdir, chat_id, message) -> dict:
     return response
 
 
+async def upload_callback(current, total, chat_id, message):
+    key = f"{chat_id}-{message.id}"
+    # if the key exists, we shouldn't send edit message
+    if not r.exists(key):
+        r.set(key, "ok", ex=EXPIRE)
+        msg = f'Uploading {round(current / total * 100, 2)}%: {current}/{total}'
+        await bot.edit_message(chat_id, message, msg)
+
+
+async def sync_edit_message(chat_id, message, msg):
+    # try to avoid flood
+    key = f"{chat_id}-{message.id}"
+    if not r.exists(key):
+        r.set(key, "ok", ex=EXPIRE)
+        await bot.edit_message(chat_id, message, msg)
+
+
 # bot starts here
 @bot.on(events.NewMessage(pattern='/start'))
 async def send_start(event):