From 3940ca9611d3d7061996e4ee35e57136d87626f9 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Sun, 18 Sep 2022 17:32:02 +0200
Subject: [PATCH] Revamp handling of partial downloads

---
 pyrogram/client.py | 22 +++++++++-------------
 1 file changed, 9 insertions(+), 13 deletions(-)

diff --git a/pyrogram/client.py b/pyrogram/client.py
index f4cee803..4ab7c344 100644
--- a/pyrogram/client.py
+++ b/pyrogram/client.py
@@ -25,7 +25,6 @@ import platform
 import re
 import shutil
 import sys
-import tempfile
 from concurrent.futures.thread import ThreadPoolExecutor
 from hashlib import sha256
 from importlib import import_module
@@ -726,32 +725,29 @@ class Client(Methods):
 
     async def handle_download(self, packet):
         file_id, directory, file_name, in_memory, file_size, progress, progress_args = packet
-        file = BytesIO() if in_memory else tempfile.NamedTemporaryFile("wb", delete=False)
 
+        os.makedirs(directory, exist_ok=True)
+        temp_file_path = os.path.abspath(re.sub("\\\\", "/", os.path.join(directory, file_name))) + ".temp"
+        file = BytesIO() if in_memory else open(temp_file_path, "wb")
+
+        # noinspection PyBroadException
         try:
             async for chunk in self.get_file(file_id, file_size, 0, 0, progress, progress_args):
                 file.write(chunk)
-        except pyrogram.StopTransmission:
+        except BaseException:
             if not in_memory:
                 file.close()
-                os.remove(file.name)
+                os.remove(temp_file_path)
 
             return None
-        except asyncio.CancelledError:
-            if not in_memory:
-                file.close()
-                os.remove(file.name)
-                
-            raise asyncio.CancelledError()
         else:
             if in_memory:
                 file.name = file_name
                 return file
             else:
-                file_path = os.path.abspath(re.sub("\\\\", "/", os.path.join(directory, file_name)))
-                os.makedirs(directory, exist_ok=True)
                 file.close()
-                shutil.move(file.name, file_path)
+                file_path = os.path.splitext(temp_file_path)[0]
+                shutil.move(temp_file_path, file_path)
                 return file_path
 
     async def get_file(