core (tools/web): helped method to download files:
this method is for internal file download, progression mechanism is not used. rel 368
This commit is contained in:
parent
303e25f72e
commit
1e1da8a7ae
1 changed files with 61 additions and 2 deletions
|
|
@ -16,13 +16,20 @@
|
|||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
from typing import Optional, Union
|
||||
from pathlib import Path
|
||||
from io import BufferedIOBase
|
||||
|
||||
from OpenSSL import SSL
|
||||
from zope.interface import implementer
|
||||
import treq
|
||||
from treq.client import HTTPClient
|
||||
from twisted.internet.interfaces import IOpenSSLClientConnectionCreator
|
||||
from twisted.internet import reactor, ssl
|
||||
from twisted.internet.interfaces import IOpenSSLClientConnectionCreator
|
||||
from twisted.web import iweb
|
||||
from twisted.web import client as http_client
|
||||
from zope.interface import implementer
|
||||
|
||||
from sat.core import exceptions
|
||||
from sat.core.log import getLogger
|
||||
|
||||
|
||||
|
|
@ -65,3 +72,55 @@ class NoCheckContextFactory:
|
|||
#: following treq doesn't check TLS, obviously it is unsecure and should not be used
|
||||
#: without explicit warning
|
||||
treq_client_no_ssl = HTTPClient(http_client.Agent(reactor, NoCheckContextFactory()))
|
||||
|
||||
|
||||
async def downloadFile(
|
||||
url: str,
|
||||
dest: Union[str, Path, BufferedIOBase],
|
||||
max_size: Optional[int] = None
|
||||
) -> None:
|
||||
"""Helper method to download a file
|
||||
|
||||
This is for internal download, for high level download with progression, use
|
||||
``plugin_misc_download``.
|
||||
|
||||
Inspired from
|
||||
https://treq.readthedocs.io/en/latest/howto.html#handling-streaming-responses
|
||||
|
||||
@param dest: destination filename or file-like object
|
||||
of it's a file-like object, you'll have to close it yourself
|
||||
@param max_size: if set, an exceptions.DataError will be raised if the downloaded file
|
||||
is bigger that given value (in bytes).
|
||||
"""
|
||||
if isinstance(dest, BufferedIOBase):
|
||||
f = dest
|
||||
must_close = False
|
||||
else:
|
||||
dest = Path(dest)
|
||||
f = dest.open("wb")
|
||||
must_close = True
|
||||
d = treq.get(url, unbuffered=True)
|
||||
written = 0
|
||||
|
||||
def write(data: bytes):
|
||||
if max_size is not None:
|
||||
nonlocal written
|
||||
written += len(data)
|
||||
if written > max_size:
|
||||
raise exceptions.DataError(
|
||||
"downloaded file is bigger than expected ({max_size})"
|
||||
)
|
||||
f.write(data)
|
||||
|
||||
d.addCallback(treq.collect, f.write)
|
||||
try:
|
||||
await d
|
||||
except exceptions.DataError as e:
|
||||
log.warning("download cancelled due to file oversized")
|
||||
raise e
|
||||
except Exception as e:
|
||||
log.error(f"Can't write file {dest}: {e}")
|
||||
raise e
|
||||
finally:
|
||||
if must_close:
|
||||
f.close()
|
||||
|
|
|
|||
Loading…
Reference in a new issue