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

Merge pull request #37 from EriHoss/flexible_media_downloads

Browse Source 
More flexible media downloads with client.download_(media/photo)
This commit is contained in:
Dan 2018-03-21 17:41:59 +01:00 committed by GitHub
commit 2122a2e1fa
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 43 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

78
pyrogram/client/client.py
View File

@ -24,7 +24,9 @@ import math
import mimetypes import mimetypes
import os import os
import re import re
import shutil
import struct import struct
import tempfile
import threading import threading
import time import time
from collections import namedtuple from collections import namedtuple
@ -503,13 +505,17 @@ class Client:
while True: while True:
media = self.download_queue.get() media = self.download_queue.get()
temp_file_path = ""
final_file_path = ""
if media is None: if media is None:
break break
try: try:
media, file_name, done, progress, path = media media, file_name, done, progress, path = media
tmp_file_name = None
directory, file_name = os.path.split(file_name)
directory = directory or "downloads"
if isinstance(media, types.MessageMediaDocument): if isinstance(media, types.MessageMediaDocument):
document = media.document document = media.document
@ -535,7 +541,7 @@ class Client:
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( temp_file_path = self.get_file(
dc_id=document.dc_id, dc_id=document.dc_id,
id=document.id, id=document.id,
access_hash=document.access_hash, access_hash=document.access_hash,
@ -558,7 +564,7 @@ class Client:
photo_loc = photo.sizes[-1].location photo_loc = photo.sizes[-1].location
tmp_file_name = self.get_file( temp_file_path = self.get_file(
dc_id=photo_loc.dc_id, dc_id=photo_loc.dc_id,
volume_id=photo_loc.volume_id, volume_id=photo_loc.volume_id,
local_id=photo_loc.local_id, local_id=photo_loc.local_id,
@ -567,30 +573,24 @@ class Client:
progress=progress progress=progress
) )
if tmp_file_name is None: if temp_file_path:
return None final_file_path = os.path.abspath(re.sub("\\\\", "/", os.path.join(directory, file_name)))
os.makedirs(directory, exist_ok=True)
if file_name is not None: shutil.move(temp_file_path, final_file_path)
path[0] = "downloads/{}".format(file_name)
try:
os.remove("downloads/{}".format(file_name))
except OSError:
pass
finally:
try:
os.renames("{}".format(tmp_file_name), "downloads/{}".format(file_name))
except OSError:
pass
except Exception as e: except Exception as e:
log.error(e, exc_info=True) log.error(e, exc_info=True)
finally:
done.set()
try: try:
os.remove("{}".format(tmp_file_name)) os.remove(temp_file_path)
except OSError: except OSError:
pass pass
else:
# TODO: "" or None for faulty download, which is better?
# os.path methods return "" in case something does not exist, I prefer this.
# For now let's keep None
path[0] = final_file_path or None
finally:
done.set()
log.debug("{} stopped".format(name)) log.debug("{} stopped".format(name))
@ -2228,9 +2228,9 @@ class Client:
version=version version=version
) )
file_name = "download_{}.temp".format(MsgId())
limit = 1024 * 1024 limit = 1024 * 1024
offset = 0 offset = 0
file_name = ""
try: try:
r = session.send( r = session.send(
@ -2242,7 +2242,9 @@ class Client:
) )
if isinstance(r, types.upload.File): if isinstance(r, types.upload.File):
with open(file_name, "wb") as f: with tempfile.NamedTemporaryFile('wb', delete=False) as f:
file_name = f.name
while True: while True:
chunk = r.bytes chunk = r.bytes
@ -2266,7 +2268,7 @@ class Client:
) )
) )
if isinstance(r, types.upload.FileCdnRedirect): elif isinstance(r, types.upload.FileCdnRedirect):
cdn_session = Session( cdn_session = Session(
r.dc_id, r.dc_id,
self.test_mode, self.test_mode,
@ -2279,7 +2281,9 @@ class Client:
cdn_session.start() cdn_session.start()
try: try:
with open(file_name, "wb") as f: with tempfile.NamedTemporaryFile('wb', delete=False) as f:
file_name = f.name
while True: while True:
r2 = cdn_session.send( r2 = cdn_session.send(
functions.upload.GetCdnFile( functions.upload.GetCdnFile(
@ -2347,6 +2351,8 @@ class Client:
os.remove(file_name) os.remove(file_name)
except OSError: except OSError:
pass pass
return ""
else: else:
return file_name return file_name
finally: finally:
@ -2605,19 +2611,20 @@ class Client:
def download_media(self, def download_media(self,
message: types.Message, message: types.Message,
file_name: str = None, file_name: str = "",
block: bool = True, block: bool = True,
progress: callable = None): progress: callable = None):
"""Use this method to download the media from a Message. """Use this method to download the media from a Message.
Files are saved in the *downloads* folder.
Args: Args:
message (:obj:`Message <pyrogram.api.types.Message>`): message (:obj:`Message <pyrogram.api.types.Message>`):
The Message containing the media. The Message containing the media.
file_name (:obj:`str`, optional): file_name (:obj:`str`, optional):
Specify a custom *file_name* to be used instead of the one provided by Telegram. A custom *file_name* to be used instead of the one provided by Telegram.
By default, all files are downloaded in the *downloads* folder in your working directory.
You can also specify a path for downloading files in a custom location: paths that end with "/"
are considered directories. All non-existent folders will be created automatically.
block (:obj:`bool`, optional): block (:obj:`bool`, optional):
Blocks the code execution until the file has been downloaded. Blocks the code execution until the file has been downloaded.
@ -2635,7 +2642,7 @@ class Client:
The size of the file. The size of the file.
Returns: Returns:
The relative path of the downloaded file. On success, the absolute path of the downloaded file as string is returned, None otherwise.
Raises: Raises:
:class:`pyrogram.Error` :class:`pyrogram.Error`
@ -2661,26 +2668,27 @@ class Client:
def download_photo(self, def download_photo(self,
photo: types.Photo or types.UserProfilePhoto or types.ChatPhoto, photo: types.Photo or types.UserProfilePhoto or types.ChatPhoto,
file_name: str = None, file_name: str = "",
block: bool = True): block: bool = True):
"""Use this method to download a photo not contained inside a Message. """Use this method to download a photo not contained inside a Message.
For example, a photo of a User or a Chat/Channel. For example, a photo of a User or a Chat/Channel.
Photos are saved in the *downloads* folder.
Args: Args:
photo (:obj:`Photo <pyrogram.api.types.Photo>` | :obj:`UserProfilePhoto <pyrogram.api.types.UserProfilePhoto>` | :obj:`ChatPhoto <pyrogram.api.types.ChatPhoto>`): photo (:obj:`Photo <pyrogram.api.types.Photo>` | :obj:`UserProfilePhoto <pyrogram.api.types.UserProfilePhoto>` | :obj:`ChatPhoto <pyrogram.api.types.ChatPhoto>`):
The photo object. The photo object.
file_name (:obj:`str`, optional): file_name (:obj:`str`, optional):
Specify a custom *file_name* to be used. A custom *file_name* to be used instead of the one provided by Telegram.
By default, all photos are downloaded in the *downloads* folder in your working directory.
You can also specify a path for downloading photos in a custom location: paths that end with "/"
are considered directories. All non-existent folders will be created automatically.
block (:obj:`bool`, optional): block (:obj:`bool`, optional):
Blocks the code execution until the photo has been downloaded. Blocks the code execution until the photo has been downloaded.
Defaults to True. Defaults to True.
Returns: Returns:
The relative path of the downloaded photo. On success, the absolute path of the downloaded photo as string is returned, None otherwise.
Raises: Raises:
:class:`pyrogram.Error` :class:`pyrogram.Error`