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:
Goffi 2022-06-29 12:07:45 +02:00
commit 1e1da8a7ae

View file

@ -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()