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

Cleaner code and some little changes

Browse Source 
TODO: "" or None for faulty download, which is better?
This commit is contained in:
Dan 2018-03-21 13:39:23 +01:00 committed by GitHub
parent 0f4e29584a
commit f6ea3e9b42
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 29 additions and 42 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

71
pyrogram/client/client.py
View File

@ -25,6 +25,7 @@ import mimetypes
import os import os
import re import re
import struct import struct
import tempfile
import threading import threading
import time import time
from collections import namedtuple from collections import namedtuple
@ -34,9 +35,6 @@ from hashlib import sha256, md5
from queue import Queue from queue import Queue
from signal import signal, SIGINT, SIGTERM, SIGABRT from signal import signal, SIGINT, SIGTERM, SIGABRT
from threading import Event, Thread from threading import Event, Thread
import tempfile
import shutil
import errno
from pyrogram.api import functions, types from pyrogram.api import functions, types
from pyrogram.api.core import Object from pyrogram.api.core import Object
@ -506,23 +504,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
download_directory = "downloads" directory, file_name = os.path.split(file_name)
directory = directory or "downloads"
if file_name.endswith('/') or file_name.endswith('\\'):
# treat the file name as a directory
download_directory = file_name
file_name = None
elif '/' in file_name or '\\' in file_name:
# use file_name as a full path instead
download_directory = ''
if isinstance(media, types.MessageMediaDocument): if isinstance(media, types.MessageMediaDocument):
document = media.document document = media.document
@ -548,7 +540,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,
@ -571,7 +563,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,
@ -580,37 +572,29 @@ class Client:
progress=progress progress=progress
) )
if tmp_file_name is None: if temp_file_path:
return None final_file_path = os.path.join(directory, file_name)
if file_name is not None:
path[0] = os.path.join(download_directory, file_name)
try:
os.remove(os.path.join(download_directory, file_name))
except OSError:
pass
finally:
try: try:
if download_directory: os.remove(final_file_path)
os.makedirs(download_directory, exist_ok=True) except OSError:
else: pass
os.makedirs(os.path.dirname(file_name), exist_ok=True)
# avoid errors moving between drives/partitions etc. os.renames(temp_file_path, final_file_path)
shutil.move(tmp_file_name, os.path.join(download_directory, file_name))
except OSError as e:
log.error(e, exc_info=True)
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(tmp_file_name) os.remove(temp_file_path)
except OSError as e: except OSError:
if e.errno != errno.ENOENT: pass
log.error(e, exc_info=True) 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))
@ -2250,7 +2234,7 @@ class Client:
limit = 1024 * 1024 limit = 1024 * 1024
offset = 0 offset = 0
file_name = None file_name = ""
try: try:
r = session.send( r = session.send(
@ -2303,6 +2287,7 @@ class Client:
try: try:
with tempfile.NamedTemporaryFile('wb', delete=False) as f: with tempfile.NamedTemporaryFile('wb', delete=False) as f:
file_name = f.name file_name = f.name
while True: while True:
r2 = cdn_session.send( r2 = cdn_session.send(
functions.upload.GetCdnFile( functions.upload.GetCdnFile(
@ -2370,6 +2355,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:
@ -2628,7 +2615,7 @@ 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.
@ -2684,7 +2671,7 @@ 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.