mir/pyrogram
2
0
Fork 0
mirror of https://github.com/pyrogram/pyrogram synced 2025-08-29 13:27:47 +00:00
Code Issues Releases Wiki Activity

Improve download_media

Browse Source
This commit is contained in:
Dan 2018-02-19 13:11:35 +01:00
parent 02da3fe947
commit 46e310722a
1 changed files with 36 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
pyrogram/client/client.py
View File

@ -276,18 +276,19 @@ class Client:
log.debug("{} started".format(name)) log.debug("{} started".format(name))
while True: while True:
message = self.download_queue.get() media = self.download_queue.get()
if message is None: if media is None:
break break
message, done = message media, file_name, done = media
try: try:
if isinstance(message.media, types.MessageMediaDocument): if isinstance(media, types.MessageMediaDocument):
document = message.media.document document = media.document
if isinstance(document, types.Document): if isinstance(document, types.Document):
if not file_name:
file_name = "doc_{}{}".format( file_name = "doc_{}{}".format(
datetime.fromtimestamp(document.date).strftime("%Y-%m-%d_%H-%M-%S"), datetime.fromtimestamp(document.date).strftime("%Y-%m-%d_%H-%M-%S"),
mimetypes.guess_extension(document.mime_type) or ".unknown" mimetypes.guess_extension(document.mime_type) or ".unknown"
@ -313,18 +314,12 @@ class Client:
version=document.version version=document.version
) )
i = 1
while True:
try: try:
os.renames("./{}".format(tmp_file_name), "./downloads/{}".format( os.remove("./downloads/{}".format(file_name))
".".join(file_name.split(".")[:-1]) except FileNotFoundError:
+ (" ({}).".format(i) if i > 1 else ".") pass
+ file_name.split(".")[-1]
)) os.renames("./{}".format(tmp_file_name), "./downloads/{}".format(file_name))
except FileExistsError:
i += 1
else:
break
done.set() done.set()
except Exception as e: except Exception as e:
@ -1940,6 +1935,9 @@ class Client:
f.flush() f.flush()
os.fsync(f.fileno()) os.fsync(f.fileno())
if len(chunk) < limit:
break
offset += limit offset += limit
except Exception as e: except Exception as e:
log.error(e) log.error(e)
@ -2301,7 +2299,10 @@ class Client:
) )
) )
def download_media(self, message: types.Message): def download_media(self, message: types.Message, file_name: str = None):
done = Event() done = Event()
self.download_queue.put((message, done)) media = message.media if isinstance(message, types.Message) else message
self.download_queue.put((media, file_name, done))
done.wait() done.wait()