# HG changeset patch
# User Goffi <goffi@goffi.org>
# Date 1656497265 -7200
# Node ID 65bac82e4049db42684474bf1a36a5db486dfcb7
# Parent  0b1c30ff2cbb9d018f0555ff702164d5e288fefe
core (tools/web): helped method to download files:

this method is for internal file download, progression mechanism is not used.

rel 368

diff -r 0b1c30ff2cbb -r 65bac82e4049 sat/tools/web.py
--- a/sat/tools/web.py	Wed Jun 29 12:06:21 2022 +0200
+++ b/sat/tools/web.py	Wed Jun 29 12:07:45 2022 +0200
@@ -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 import reactor, ssl
 from twisted.internet.interfaces import IOpenSSLClientConnectionCreator
-from twisted.internet import reactor, ssl
 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 @@
 #: 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()