Download files in chunks

Author: hovnatan

superannotate/db/exports.py

@@ -6,6 +6,7 @@
 import zipfile
 from datetime import datetime
 from pathlib import Path
+import shutil
 
 import boto3
 import requests
@@ -209,6 +210,15 @@ def __upload_files_to_aws_thread(
             already_uploaded[i] = True
 
 
+def _download_file(url, local_filename):
+    with requests.get(url, stream=True) as r:
+        r.raise_for_status()
+        with open(local_filename, 'wb') as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                f.write(chunk)
+    return local_filename
+
+
 def download_export(
     project, export, folder_path, extract_zip_contents=True, to_s3_bucket=None
 ):
@@ -243,25 +253,26 @@ def download_export(
         break
 
     filename = Path(res['path']).name
-    r = requests.get(res['download'], allow_redirects=True)
-    if to_s3_bucket is None:
-        filepath = Path(folder_path) / filename
-        open(filepath, 'wb').write(r.content)
-        if extract_zip_contents:
-            with zipfile.ZipFile(filepath, 'r') as f:
-                f.extractall(folder_path)
-            Path.unlink(filepath)
-            logger.info("Extracted %s to folder %s", filepath, folder_path)
-        else:
-            logger.info("Downloaded export ID %s to %s", res['id'], filepath)
-    else:
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            filepath = Path(tmpdirname) / filename
-            open(filepath, 'wb').write(r.content)
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        temp_filepath = Path(tmpdirname) / filename
+        _download_file(res['download'], temp_filepath)
+        if to_s3_bucket is None:
+            filepath = Path(folder_path) / filename
+            shutil.copyfile(temp_filepath, filepath)
             if extract_zip_contents:
                 with zipfile.ZipFile(filepath, 'r') as f:
-                    f.extractall(tmpdirname)
+                    f.extractall(folder_path)
                 Path.unlink(filepath)
+                logger.info("Extracted %s to folder %s", filepath, folder_path)
+            else:
+                logger.info(
+                    "Downloaded export ID %s to %s", res['id'], filepath
+                )
+        else:
+            if extract_zip_contents:
+                with zipfile.ZipFile(temp_filepath, 'r') as f:
+                    f.extractall(tmpdirname)
+                Path.unlink(temp_filepath)
             files_to_upload = []
             for file in Path(tmpdirname).rglob("*.*"):
                 if not file.is_file():
@@ -296,4 +307,4 @@ def download_export(
                     t.join()
             finish_event.set()
             tqdm_thread.join()
-        logger.info("Exported to AWS %s/%s", to_s3_bucket, folder_path)
+            logger.info("Exported to AWS %s/%s", to_s3_bucket, folder_path)