X Tutup
Skip to content

Commit f6332d4

Browse files
tsnoamtsnoam
authored
Improved File.download (python-telegram-bot#1019)
 - File.download_as_bytearray - new method to get a d/led file as bytearray This is much more convenient and straight forward than using a file object. - File.download(): Now returns a meaningful return value - File.download*(): New and/or better unitests
1 parent b275031 commit f6332d4

File tree

2 files changed

+92
-12
lines changed

2 files changed

+92
-12
lines changed

‎telegram/files/file.py‎

Lines changed: 34 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -74,36 +74,62 @@ def download(self, custom_path=None, out=None, timeout=None):
7474
that object using the ``out.write`` method.
7575
7676
Note:
77-
`custom_path` and `out` are mutually exclusive.
77+
:attr:`custom_path` and :attr:`out` are mutually exclusive.
7878
7979
Args:
8080
custom_path (:obj:`str`, optional): Custom path.
81-
out (:obj:`object`, optional): A file-like object. Must be opened in binary mode, if
82-
applicable.
81+
out (:obj:`io.BufferedWriter`, optional): A file-like object. Must be opened for
82+
writing in binary mode, if applicable.
8383
timeout (:obj:`int` | :obj:`float`, optional): If this value is specified, use it as
8484
the read timeout from the server (instead of the one specified during creation of
8585
the connection pool).
8686
87+
Returns:
88+
:obj:`str` | :obj:`io.BufferedWriter`: The same object as :attr:`out` if specified.
89+
Otherwise, returns the filename downloaded to.
90+
8791
Raises:
88-
ValueError: If both ``custom_path`` and ``out`` are passed.
92+
ValueError: If both :attr:`custom_path` and :attr:`out` are passed.
8993
9094
"""
9195
if custom_path is not None and out is not None:
9296
raise ValueError('custom_path and out are mutually exclusive')
9397

9498
# Convert any UTF-8 char into a url encoded ASCII string.
95-
sres = urllib_parse.urlsplit(self.file_path)
96-
url = urllib_parse.urlunsplit(urllib_parse.SplitResult(
97-
sres.scheme, sres.netloc, urllib_parse.quote(sres.path), sres.query, sres.fragment))
99+
url = self._get_encoded_url()
98100

99101
if out:
100102
buf = self.bot.request.retrieve(url)
101103
out.write(buf)
102-
104+
return out
103105
else:
104106
if custom_path:
105107
filename = custom_path
106108
else:
107109
filename = basename(self.file_path)
108110

109111
self.bot.request.download(url, filename, timeout=timeout)
112+
return filename
113+
114+
def _get_encoded_url(self):
115+
"""Convert any UTF-8 char in :obj:`File.file_path` into a url encoded ASCII string."""
116+
sres = urllib_parse.urlsplit(self.file_path)
117+
return urllib_parse.urlunsplit(urllib_parse.SplitResult(
118+
sres.scheme, sres.netloc, urllib_parse.quote(sres.path), sres.query, sres.fragment))
119+
120+
def download_as_bytearray(self, buf=None):
121+
"""Download this file and return it as a bytearray.
122+
123+
Args:
124+
buf (:obj:`bytearray`, optional): Extend the given bytearray with the downloaded data.
125+
126+
Returns:
127+
:obj:`bytearray`: The same object as :attr:`buf` if it was specified. Otherwise a newly
128+
allocated :obj:`bytearray`.
129+
130+
"""
131+
if buf is None:
132+
buf = bytearray()
133+
134+
buf.extend(self.bot.request.retrieve(self._get_encoded_url()))
135+
return buf

‎tests/test_file.py‎

Lines changed: 58 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#!/usr/bin/env python
2+
# -*- coding: utf-8 -*-
23
#
34
# A library that provides a Python interface to the Telegram Bot API
45
# Copyright (C) 2015-2018
@@ -16,6 +17,8 @@
1617
#
1718
# You should have received a copy of the GNU Lesser Public License
1819
# along with this program. If not, see [http://www.gnu.org/licenses/].
20+
import os
21+
from tempfile import TemporaryFile, mkstemp
1922

2023
import pytest
2124
from flaky import flaky
@@ -36,6 +39,7 @@ class TestFile(object):
3639
file_path = (
3740
u'https://api.org/file/bot133505823:AAHZFMHno3mzVLErU5b5jJvaeG--qUyLyG0/document/file_3')
3841
file_size = 28232
42+
file_content = u'Saint-Saëns'.encode('utf-8') # Intentionally contains unicode chars.
3943

4044
def test_de_json(self, bot):
4145
json_dict = {
@@ -65,11 +69,61 @@ def test_error_get_empty_file_id(self, bot):
6569

6670
def test_download(self, monkeypatch, file):
6771
def test(*args, **kwargs):
68-
raise TelegramError('test worked')
72+
return self.file_content
6973

70-
monkeypatch.setattr('telegram.utils.request.Request.download', test)
71-
with pytest.raises(TelegramError, match='test worked'):
72-
file.download()
74+
monkeypatch.setattr('telegram.utils.request.Request.retrieve', test)
75+
out_file = file.download()
76+
77+
try:
78+
with open(out_file, 'rb') as fobj:
79+
assert fobj.read() == self.file_content
80+
finally:
81+
os.unlink(out_file)
82+
83+
def test_download_custom_path(self, monkeypatch, file):
84+
def test(*args, **kwargs):
85+
return self.file_content
86+
87+
monkeypatch.setattr('telegram.utils.request.Request.retrieve', test)
88+
file_handle, custom_path = mkstemp()
89+
try:
90+
out_file = file.download(custom_path)
91+
assert out_file == custom_path
92+
93+
with open(out_file, 'rb') as fobj:
94+
assert fobj.read() == self.file_content
95+
finally:
96+
os.close(file_handle)
97+
os.unlink(custom_path)
98+
99+
def test_download_file_obj(self, monkeypatch, file):
100+
def test(*args, **kwargs):
101+
return self.file_content
102+
103+
monkeypatch.setattr('telegram.utils.request.Request.retrieve', test)
104+
with TemporaryFile() as custom_fobj:
105+
out_fobj = file.download(out=custom_fobj)
106+
assert out_fobj is custom_fobj
107+
108+
out_fobj.seek(0)
109+
assert out_fobj.read() == self.file_content
110+
111+
def test_download_bytearray(self, monkeypatch, file):
112+
def test(*args, **kwargs):
113+
return self.file_content
114+
115+
monkeypatch.setattr('telegram.utils.request.Request.retrieve', test)
116+
117+
# Check that a download to a newly allocated bytearray works.
118+
buf = file.download_as_bytearray()
119+
assert buf == bytearray(self.file_content)
120+
121+
# Check that a download to a given bytearray works (extends the bytearray).
122+
buf2 = buf[:]
123+
buf3 = file.download_as_bytearray(buf=buf2)
124+
assert buf3 is buf2
125+
assert buf2[len(buf):] == buf
126+
assert buf2[:len(buf)] == buf
73127

74128
def test_equality(self, bot):
75129
a = File(self.file_id, bot)

0 commit comments

Comments
 (0)
X Tutup