Support proxy in download_folder

wkentaro · Feb 21, 2023 · 311d5c0 · 311d5c0
1 parent e4cb896
commit 311d5c0
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 17 deletions.
diff --git a/gdown/download.py b/gdown/download.py
@@ -61,9 +61,17 @@ def get_url_from_gdrive_confirmation(contents):
     return url
 
 
-def _get_session(use_cookies, return_cookies_file=False):
+def _get_session(proxy, use_cookies, return_cookies_file=False):
     sess = requests.session()
 
+    sess.headers.update(
+        {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6)"}
+    )
+
+    if proxy is not None:
+        sess.proxies = {"http": proxy, "https": proxy}
+        print("Using proxy:", proxy, file=sys.stderr)
+
     # Load cookies if exists
     cookies_file = osp.join(home, ".cache/gdown/cookies.json")
     if osp.exists(cookies_file) and use_cookies:
@@ -131,13 +139,9 @@ def download(
     url_origin = url
 
     sess, cookies_file = _get_session(
-        use_cookies=use_cookies, return_cookies_file=True
+        proxy=proxy, use_cookies=use_cookies, return_cookies_file=True
     )
 
-    if proxy is not None:
-        sess.proxies = {"http": proxy, "https": proxy}
-        print("Using proxy:", proxy, file=sys.stderr)
-
     gdrive_file_id, is_gdrive_download_link = parse_url(url, warning=not fuzzy)
 
     if fuzzy and gdrive_file_id:
@@ -146,15 +150,15 @@ def download(
         url_origin = url
         is_gdrive_download_link = True
 
-    headers = {
-        "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_10_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36"  # NOQA
-    }
-
     while True:
         try:
-            res = sess.get(url, headers=headers, stream=True, verify=verify)
+            res = sess.get(url, stream=True, verify=verify)
         except requests.exceptions.ProxyError as e:
-            print("An error has occurred using proxy:", proxy, file=sys.stderr)
+            print(
+                "An error has occurred using proxy:",
+                sess.proxies,
+                file=sys.stderr,
+            )
             print(e, file=sys.stderr)
             return
 
@@ -246,7 +250,7 @@ def download(
         f = output
 
     if tmp_file is not None and f.tell() != 0:
-        headers["Range"] = "bytes={}-".format(f.tell())
+        headers = {"Range": "bytes={}-".format(f.tell())}
         res = sess.get(url, headers=headers, stream=True, verify=verify)
 
     if not quiet:

diff --git a/gdown/download_folder.py b/gdown/download_folder.py
@@ -12,6 +12,7 @@
 import warnings
 
 import bs4
+import requests
 
 from .download import _get_session
 from .download import download
@@ -116,14 +117,21 @@ def _download_and_parse_google_drive_link(
     else:
         url += "?hl=en"
 
-    folder_page = sess.get(url, verify=verify)
+    try:
+        res = sess.get(url, verify=verify)
+    except requests.exceptions.ProxyError as e:
+        print(
+            "An error has occurred using proxy:", sess.proxies, file=sys.stderr
+        )
+        print(e, file=sys.stderr)
+        return False, None
 
-    if folder_page.status_code != 200:
+    if res.status_code != 200:
         return False, None
 
     gdrive_file, id_name_type_iter = _parse_google_drive_file(
         url=url,
-        content=folder_page.text,
+        content=res.text,
     )
 
     for child_id, child_name, child_type in id_name_type_iter:
@@ -249,7 +257,7 @@ def download_folder(
     if id is not None:
         url = "https://drive.google.com/drive/folders/{id}".format(id=id)
 
-    sess = _get_session(use_cookies=use_cookies)
+    sess = _get_session(proxy=proxy, use_cookies=use_cookies)
 
     if not quiet:
         print("Retrieving folder list", file=sys.stderr)