Xtao-Labs/MTPyroger
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add support for downloading files to file pointer, fix for https://github.com/pyrogram/pyrogram/issues/284

Browse Source
This commit is contained in:
Octo 2019-07-20 09:20:57 +03:00 committed by Dan
parent 55d0b93cf0
commit 1e8c9812a1
3 changed files with 120 additions and 93 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
pyrogram/client/client.py
View File

@ -15,7 +15,7 @@
#
# You should have received a copy of the GNU Lesser General Public License
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
import io
import logging
import math
import os
@ -1231,9 +1231,9 @@ class Client(Methods, BaseClient):
temp_file_path = ""
final_file_path = ""
path = [None]
try:
data, directory, file_name, done, progress, progress_args, path = packet
data, done, progress, progress_args, out, path, to_file = packet
temp_file_path = self.get_file(
media_type=data.media_type,
@ -1250,13 +1250,15 @@ class Client(Methods, BaseClient):
file_size=data.file_size,
is_big=data.is_big,
progress=progress,
progress_args=progress_args
progress_args=progress_args,
out=out
)
if temp_file_path:
final_file_path = os.path.abspath(re.sub("\\\\", "/", os.path.join(directory, file_name)))
os.makedirs(directory, exist_ok=True)
shutil.move(temp_file_path, final_file_path)
if to_file:
final_file_path = out.name
else:
final_file_path = ''
if to_file:
out.close()
except Exception as e:
log.error(e, exc_info=True)
@ -1864,7 +1866,8 @@ class Client(Methods, BaseClient):
file_size: int,
is_big: bool,
progress: callable,
progress_args: tuple = ()
progress_args: tuple = (),
out: io.IOBase = None
) -> str:
with self.media_sessions_lock:
session = self.media_sessions.get(dc_id, None)
@ -1950,7 +1953,10 @@ class Client(Methods, BaseClient):
limit = 1024 * 1024
offset = 0
file_name = ""
if not out:
f = tempfile.NamedTemporaryFile("wb", delete=False)
else:
f = out
try:
r = session.send(
functions.upload.GetFile(
@ -1961,7 +1967,7 @@ class Client(Methods, BaseClient):
)
if isinstance(r, types.upload.File):
with tempfile.NamedTemporaryFile("wb", delete=False) as f:
if hasattr(f, "name"):
file_name = f.name
while True:
@ -1976,6 +1982,7 @@ class Client(Methods, BaseClient):
if progress:
progress(
min(offset, file_size)
if file_size != 0
else offset,
@ -2003,7 +2010,7 @@ class Client(Methods, BaseClient):
self.media_sessions[r.dc_id] = cdn_session
try:
with tempfile.NamedTemporaryFile("wb", delete=False) as f:
if hasattr(f, "name"):
file_name = f.name
while True:
@ -2058,6 +2065,7 @@ class Client(Methods, BaseClient):
if progress:
progress(
min(offset, file_size)
if file_size != 0
else offset,
@ -2074,6 +2082,7 @@ class Client(Methods, BaseClient):
log.error(e, exc_info=True)
try:
if out:
os.remove(file_name)
except OSError:
pass

13
pyrogram/client/methods/messages/download_media.py
View File

@ -17,7 +17,9 @@
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
import binascii
import io
import os
import re
import struct
import time
from datetime import datetime
@ -37,6 +39,7 @@ class DownloadMedia(BaseClient):
message: Union["pyrogram.Message", str],
file_ref: str = None,
file_name: str = DEFAULT_DOWNLOAD_DIR,
out: io.IOBase = None,
block: bool = True,
progress: callable = None,
progress_args: tuple = ()
@ -58,6 +61,9 @@ class DownloadMedia(BaseClient):
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.
out (``io.IOBase``, *optional*):
A custom *file-like object* to be used when downloading file. Overrides file_name
block (``bool``, *optional*):
Blocks the code execution until the file has been downloaded.
Defaults to True.
@ -238,6 +244,13 @@ class DownloadMedia(BaseClient):
extension
)
if not out:
out = open(os.path.abspath(re.sub("\\\\", "/", os.path.join(directory, file_name))), 'wb')
os.makedirs(directory, exist_ok=True)
to_file = True
else:
to_file = False
self.download_queue.put((data, done, progress, progress_args, out, path, to_file))
# Cast to string because Path objects aren't supported by Python 3.5
self.download_queue.put((data, str(directory), str(file_name), done, progress, progress_args, path))

7
pyrogram/client/types/messages_and_media/message.py
View File

@ -15,7 +15,7 @@
#
# You should have received a copy of the GNU Lesser General Public License
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
import io
from functools import partial
from typing import List, Match, Union
@ -2985,6 +2985,7 @@ class Message(Object, Update):
def download(
self,
file_name: str = "",
out: io.IOBase = None,
block: bool = True,
progress: callable = None,
progress_args: tuple = ()
@ -3009,6 +3010,9 @@ class Message(Object, Update):
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.
out (``io.IOBase``, *optional*):
A custom *file-like object* to be used when downloading file. Overrides file_name
block (``bool``, *optional*):
Blocks the code execution until the file has been downloaded.
Defaults to True.
@ -3045,6 +3049,7 @@ class Message(Object, Update):
return self._client.download_media(
message=self,
file_name=file_name,
out=out,
block=block,
progress=progress,
progress_args=progress_args,