mir/pyrogram
2
0
Fork 0
mirror of https://github.com/pyrogram/pyrogram synced 2025-08-28 21:07:59 +00:00
Code Issues Releases Wiki Activity

Revamp handling of partial downloads

Browse Source
This commit is contained in:
Dan 2022-09-18 17:32:02 +02:00
parent 468ebf50cf
commit 3940ca9611
1 changed files with 9 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
pyrogram/client.py
View File

@ -25,7 +25,6 @@ import platform
import re import re
import shutil import shutil
import sys import sys
import tempfile
from concurrent.futures.thread import ThreadPoolExecutor from concurrent.futures.thread import ThreadPoolExecutor
from hashlib import sha256 from hashlib import sha256
from importlib import import_module from importlib import import_module
@ -726,32 +725,29 @@ class Client(Methods):
async def handle_download(self, packet): async def handle_download(self, packet):
file_id, directory, file_name, in_memory, file_size, progress, progress_args = 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: try:
async for chunk in self.get_file(file_id, file_size, 0, 0, progress, progress_args): async for chunk in self.get_file(file_id, file_size, 0, 0, progress, progress_args):
file.write(chunk) file.write(chunk)
except pyrogram.StopTransmission: except BaseException:
if not in_memory: if not in_memory:
file.close() file.close()
os.remove(file.name) os.remove(temp_file_path)
return None return None
except asyncio.CancelledError:
if not in_memory:
file.close()
os.remove(file.name)
raise asyncio.CancelledError()
else: else:
if in_memory: if in_memory:
file.name = file_name file.name = file_name
return file return file
else: else:
file_path = os.path.abspath(re.sub("\\\\", "/", os.path.join(directory, file_name)))
os.makedirs(directory, exist_ok=True)
file.close() 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 return file_path
async def get_file( async def get_file(