From f0c7a600813b0f3f20bf6a3dbcbb0fb5a1815d2b Mon Sep 17 00:00:00 2001
From: BennyThink <benny.think@gmail.com>
Date: Wed, 9 Feb 2022 11:35:18 +0800
Subject: [PATCH] audio only format

fix #59
---
 ytdlbot/downloader.py | 29 ++++++++++++++++++++++-------
 ytdlbot/limit.py      |  3 ---
 ytdlbot/tasks.py      | 10 ++++++----
 ytdlbot/ytdl_bot.py   |  2 +-
 4 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/ytdlbot/downloader.py b/ytdlbot/downloader.py
index 45cfc91f..d1b114b8 100644
--- a/ytdlbot/downloader.py
+++ b/ytdlbot/downloader.py
@@ -24,7 +24,7 @@
 from tqdm import tqdm
 from yt_dlp import DownloadError
 
-from config import ENABLE_VIP, MAX_DURATION, TG_MAX_SIZE
+from config import AUDIO_FORMAT, ENABLE_VIP, MAX_DURATION, TG_MAX_SIZE
 from db import Redis
 from limit import VIP
 from utils import (adjust_formats, apply_log_formatter, current_time,
@@ -139,10 +139,9 @@ def convert_to_mp4(resp: dict, bot_msg):
                         bot_msg.chat.id,
                         "You're not VIP, so you can't convert longer video to streaming formats.")
                     break
-                pobj = pathlib.Path(path)
-                edit_text(bot_msg, f"{current_time()}: Converting {pobj.name} to mp4. Please wait.")
-                new_file_path = pobj.with_suffix(".mp4")
-                cmd = ["ffmpeg", "-i", path, new_file_path]
+                edit_text(bot_msg, f"{current_time()}: Converting {path.name} to mp4. Please wait.")
+                new_file_path = path.with_suffix(".mp4")
+                cmd = ["ffmpeg", "-y", "-i", path, new_file_path]
                 logging.info("Detected %s, converting to mp4...", mime)
                 subprocess.check_output(cmd)
                 index = resp["filepath"].index(path)
@@ -183,7 +182,6 @@ def ytdl_download(url, tempdir, bm) -> dict:
     ]
     adjust_formats(chat_id, url, formats)
     add_instagram_cookies(url, ydl_opts)
-    # TODO it appears twitter download on macOS will fail. Don't know why...Linux's fine.
     for f in formats:
         if f:
             ydl_opts["format"] = f
@@ -207,7 +205,7 @@ def ytdl_download(url, tempdir, bm) -> dict:
         return response
 
     for i in os.listdir(tempdir):
-        p: "str" = os.path.join(tempdir, i)
+        p = pathlib.Path(tempdir, i)
         file_size = os.stat(p).st_size
         if ENABLE_VIP:
             remain, _, ttl = VIP().check_remaining_quota(chat_id)
@@ -227,11 +225,28 @@ def ytdl_download(url, tempdir, bm) -> dict:
     if settings[2] == "video" or isinstance(settings[2], MagicMock):
         # only convert if send type is video
         convert_to_mp4(response, bm)
+    if settings[2] == "audio":
+        check_audio_format(response)
     # disable it for now
     # split_large_video(response)
     return response
 
 
+def check_audio_format(resp: "dict"):
+    if resp["status"]:
+        # all_converted = []
+        path: pathlib.PosixPath
+        for path in resp["filepath"]:
+            # if we can't guess file type, we assume it's video/mp4
+            if path.suffix != f".{AUDIO_FORMAT}":
+                new_path = path.with_suffix(f".{AUDIO_FORMAT}")
+                cmd = 'ffmpeg -y -i "{}" "{}"'.format(path, new_path)
+                subprocess.check_output(cmd, shell=True)
+                path.unlink()
+                index = resp["filepath"].index(path)
+                resp["filepath"][index] = new_path
+
+
 def add_instagram_cookies(url: "str", opt: "dict"):
     if url.startswith("https://www.instagram.com"):
         opt["cookiefi22"] = pathlib.Path(__file__).parent.joinpath("instagram.com_cookies.txt").as_posix()
diff --git a/ytdlbot/limit.py b/ytdlbot/limit.py
index 3d74b4ff..a7edbe36 100644
--- a/ytdlbot/limit.py
+++ b/ytdlbot/limit.py
@@ -295,6 +295,3 @@ def subscribe_query():
             print(f"{has} - {uid}")
 
 
-if __name__ == '__main__':
-    a = VIP.extract_canonical_link("https://www.youtube.com/shorts/YrnvPPGznXM")
-    print(a)
diff --git a/ytdlbot/tasks.py b/ytdlbot/tasks.py
index 56ed2097..0458f8de 100644
--- a/ytdlbot/tasks.py
+++ b/ytdlbot/tasks.py
@@ -12,12 +12,12 @@
 import pathlib
 import re
 import subprocess
-import psutil
 import tempfile
 import threading
 import time
 from urllib.parse import quote_plus
 
+import psutil
 import requests
 from apscheduler.schedulers.background import BackgroundScheduler
 from celery import Celery
@@ -223,9 +223,11 @@ def normal_audio(bot_msg, client):
         Redis().update_metrics("audio_success")
 
 
-def get_worker_status():
+def get_dl_source():
     worker_name = os.getenv("WORKER_NAME")
-    return f"Downloaded by  {worker_name}"
+    if worker_name:
+        return f"Downloaded by  {worker_name}"
+    return ""
 
 
 def upload_transfer_sh(video_paths) -> "str":
@@ -273,7 +275,7 @@ def ytdl_normal_download(bot_msg, client, url):
                 return
 
             meta = get_metadata(video_path)
-            worker = "Downloaded by {}".format(os.getenv("WORKER_NAME", "Unknown"))
+            worker = get_dl_source()
             cap = f"`{filename}`\n\n{url}\n\nInfo: {meta['width']}x{meta['height']} {size} {meta['duration']}s" \
                   f"\n{remain}\n{worker}"
             settings = get_user_settings(str(chat_id))
diff --git a/ytdlbot/ytdl_bot.py b/ytdlbot/ytdl_bot.py
index 6a2ceb46..512ff9ea 100644
--- a/ytdlbot/ytdl_bot.py
+++ b/ytdlbot/ytdl_bot.py
@@ -249,7 +249,7 @@ def download_handler(client: "Client", message: "types.Message"):
 
     red.update_metrics("video_request")
     text = bot_text.get_receive_link_text()
-    time.sleep(random.random() * 2)
+    time.sleep(random.random() / 2)
     try:
         # raise pyrogram.errors.exceptions.FloodWait(10)
         bot_msg: typing.Union["types.Message", "typing.Any"] = message.reply_text(text, quote=True)