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

71
pyrogram/client/client.py
View File

@ -276,35 +276,36 @@ 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):
file_name = "doc_{}{}".format( if not file_name:
datetime.fromtimestamp(document.date).strftime("%Y-%m-%d_%H-%M-%S"), file_name = "doc_{}{}".format(
mimetypes.guess_extension(document.mime_type) or ".unknown" datetime.fromtimestamp(document.date).strftime("%Y-%m-%d_%H-%M-%S"),
) mimetypes.guess_extension(document.mime_type) or ".unknown"
)
for i in document.attributes: for i in document.attributes:
if isinstance(i, types.DocumentAttributeFilename): if isinstance(i, types.DocumentAttributeFilename):
file_name = i.file_name file_name = i.file_name
break break
elif isinstance(i, types.DocumentAttributeSticker): elif isinstance(i, types.DocumentAttributeSticker):
file_name = file_name.replace("doc", "sticker") file_name = file_name.replace("doc", "sticker")
elif isinstance(i, types.DocumentAttributeAudio): elif isinstance(i, types.DocumentAttributeAudio):
file_name = file_name.replace("doc", "audio") file_name = file_name.replace("doc", "audio")
elif isinstance(i, types.DocumentAttributeVideo): elif isinstance(i, types.DocumentAttributeVideo):
file_name = file_name.replace("doc", "video") file_name = file_name.replace("doc", "video")
elif isinstance(i, types.DocumentAttributeAnimated): elif isinstance(i, types.DocumentAttributeAnimated):
file_name = file_name.replace("doc", "gif") file_name = file_name.replace("doc", "gif")
tmp_file_name = self.get_file( tmp_file_name = self.get_file(
dc_id=document.dc_id, dc_id=document.dc_id,
@ -313,18 +314,12 @@ class Client:
version=document.version version=document.version
) )
i = 1 try:
while True: os.remove("./downloads/{}".format(file_name))
try: except FileNotFoundError:
os.renames("./{}".format(tmp_file_name), "./downloads/{}".format( pass
".".join(file_name.split(".")[:-1])
+ (" ({}).".format(i) if i > 1 else ".") os.renames("./{}".format(tmp_file_name), "./downloads/{}".format(file_name))
+ file_name.split(".")[-1]
))
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()