Merge pull request #1141 from nextgenusfs/bugfix_1138

fix bug #1138 by replacing how url requests are made

Author: hyphaltip

funannotate/setupDB.py

@@ -13,9 +13,7 @@
 import datetime
 import requests
 import hashlib
-import socket
 import shutil
-import errno
 import json
 import xml.etree.cElementTree as cElementTree
 import funannotate.library as lib
@@ -72,32 +70,28 @@ def download(url, name, wget=False):
     else:
         file_name = name
         try:
-            u = urlopen(url)
-            f = open(file_name, 'wb')
-            meta = u.info()
-            file_size = 0
-            for x in meta.items():
-                if x[0].lower() == 'content-length':
-                    file_size = int(x[1])
-            lib.log.info("Downloading: {0} Bytes: {1}".format(url, file_size))
-            file_size_dl = 0
-            block_sz = 8192
-            while True:
-                buffer = u.read(block_sz)
-                if not buffer:
-                    break
-                file_size_dl += len(buffer)
-                f.write(buffer)
-                p = float(file_size_dl) / file_size
-                status = r"{0}  [{1:.2%}]".format(file_size_dl, p)
-                status = status + chr(8)*(len(status)+1)
-                sys.stdout.write(status)
-            sys.stdout.flush()
-            f.close()
-        except socket.error as e:
-            if e.errno != errno.ECONNRESET:
-                raise
-            pass
+            # fix bug 1138
+            # Use requests instead of urlopen to handle 308 redirects - was necessary for osf.io requests
+            with requests.get(url, stream=True, allow_redirects=True) as r:
+                r.raise_for_status()
+                file_size = int(r.headers.get('content-length', 0))
+                lib.log.info("Downloading:  {0} Bytes:  {1}".format(url, file_size))
+                with open(file_name, 'wb') as f:
+                    file_size_dl = 0
+                    block_sz = 8192
+                    for chunk in r.iter_content(chunk_size=block_sz):
+                        if chunk:
+                            file_size_dl += len(chunk)
+                            f.write(chunk)
+                            if file_size > 0:
+                                p = float(file_size_dl) / file_size
+                                status = r"{0} [{1:.2%}]".format(file_size_dl, p)
+                                status = status + chr(8)*(len(status)+1)
+                                sys.stdout.write(status)
+                sys.stdout.flush()
+        except requests.exceptions.RequestException as e:
+            lib.log.error("Download failed: {}".format(e))
+            raise
 
 
 def download_busco_v5(dest, taxa=['Eukaryota'], wget=False):