From 32be202a4b712a6cf5ec481d71f5a8b53f88e624 Mon Sep 17 00:00:00 2001
From: Sanuja Seneviratne <66342986+SanujaNS@users.noreply.github.com>
Date: Tue, 23 Jul 2024 18:14:17 +0530
Subject: [PATCH] improve /direct command (#410)

---
 ytdlbot/tasks.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/ytdlbot/tasks.py b/ytdlbot/tasks.py
index 4e0b192a..1554ad6b 100644
--- a/ytdlbot/tasks.py
+++ b/ytdlbot/tasks.py
@@ -298,6 +298,12 @@ def direct_normal_download(client: Client, bot_msg: typing.Union[types.Message,
             downloaded += len(chunk)
         logging.info("Downloaded file %s", filename)
         st_size = os.stat(filepath).st_size
+        ext = filetype.guess_extension(filepath)
+        # Rename file if it doesn't have extension
+        if ext is not None and not filepath.endswith(ext):
+            new_filename = f"{filepath}.{ext}"
+            os.rename(filepath, new_filename)
+            filepath = new_filename
 
         client.send_chat_action(chat_id, enums.ChatAction.UPLOAD_DOCUMENT)
         client.send_document(
@@ -359,10 +365,9 @@ def leech_normal_download(client: Client, bot_msg: typing.Union[types.Message, t
     bot_msg.edit_text(f"Download Complete", disable_web_page_preview=True)
     ext = filetype.guess_extension(file_path_obj)
     # Rename file if it doesn't have extension
-    if ext is not None:
-        if not filename.endswith(ext):
-            new_filename = f"{tempdir}/{filename}.{ext}"
-            os.rename(file_path_obj, new_filename)
+    if ext is not None and not filename.endswith(ext):
+        new_filename = f"{tempdir}/{filename}.{ext}"
+        os.rename(file_path_obj, new_filename)
     # Get file path of the downloaded file to upload
     video_paths = list(pathlib.Path(tempdir).glob("*"))
     client.send_chat_action(chat_id, enums.ChatAction.UPLOAD_DOCUMENT)