added resume progress bar

pyatmos/msise/spaceweather.py

@@ -19,8 +19,6 @@ from datetime import datetime,timedelta
 from os import path,makedirs,remove
 from pathlib import Path
 from urllib.request import urlretrieve
-from colorama import Fore
-from .tqdmupto import TqdmUpTo
 
 # =================== download and update sw data  =================== #
 
@@ -58,21 +56,15 @@ def download_sw(direc=None):
     url = 'https://www.celestrak.com/SpaceData/SW-All.txt'
 
     if not path.exists(direc): makedirs(direc)
-    bar_format = "{l_bar}%s{bar}%s{r_bar}" % (Fore.BLUE, Fore.RESET)
     if not path.exists(swfile):
         desc = 'Downloading the latest space weather data from CELESTRAK'
-        with TqdmUpTo(unit='B', unit_scale=True, desc=desc,bar_format = bar_format,position=0) as t:
-            urlretrieve(url, swfile,reporthook=t.update_to)
-            t.total = t.n
-
+        tqdm_request(url,direc,'SW-All.txt',desc)
     else:
         modified_time = datetime.fromtimestamp(path.getmtime(swfile))
         if datetime.now() > modified_time + timedelta(days=1):
             remove(swfile)
             desc = 'Updating the space weather data from CELESTRAK'
-            with TqdmUpTo(unit='B', unit_scale=True, desc=desc,bar_format = bar_format,position=0) as t:
-                urlretrieve(url, swfile,reporthook=t.update_to)    
-                t.total = t.n
+            tqdm_request(url,direc,'SW-All.txt',desc)   
         else:
             print('The space weather data in {:s} is already the latest.'.format(direc))   
     return swfile

pyatmos/msise/tqdmupto.py

@@ -1,19 +0,0 @@
-from tqdm import tqdm
-
-class TqdmUpTo(tqdm):
-    '''
-    Provides `update_to(n)` which uses `tqdm.update(delta_n)`.
-    For more details, please refer to https://pypi.org/project/tqdm/ 
-    '''
-    def update_to(self, b=1, bsize=1, tsize=None):
-        '''
-        b  : int, optional
-            Number of blocks transferred so far [default: 1].
-        bsize  : int, optional
-            Size of each block (in tqdm units) [default: 1].
-        tsize  : int, optional
-            Total size (in tqdm units). If [default: None] remains unchanged.
-        '''
-        if tsize is not None:
-            self.total = tsize
-        self.update(b * bsize - self.n)  # will also set self.n = b * bsize

pyatmos/msise/utils.py

@@ -1,3 +1,6 @@
+import requests
+from tqdm import tqdm
+from colorama import Fore
 # ------------------------------------------------------------------- #
 # ----------------------------- utilities --------------------------- #
 # ------------------------------------------------------------------- #
@@ -31,3 +34,27 @@ def hms2s(h,m,s):
 def hms2h(h,m,s):
     return h + m/60 + s/3.6E3
 
+def tqdm_request(url,dir_to,file,desc):
+    block_size = 1024
+    bar_format = "{l_bar}%s{bar}%s{r_bar}" % (Fore.BLUE, Fore.RESET)
+    for idownload in range(5):
+        try:
+            local_file = open(dir_to + file, 'ab')
+            pos = local_file.tell()
+            resume_header = {'Range': f'bytes={pos}-'}
+            res = requests.get(url,stream=True,timeout=100,headers=resume_header)
+            total_size = int(res.headers.get('content-length', 0))
+            pbar = tqdm(desc = desc,total=total_size,unit='B',unit_scale=True,bar_format = bar_format,position=0,initial=pos)
+            for chunk in res.iter_content(block_size):
+                pbar.update(len(chunk))
+                local_file.write(chunk)  
+            break
+        except: 
+            sleep(2)
+            if idownload == 4:
+                remove(dir_to + file)
+                print('No response, skip this file.') 
+        finally:
+            pbar.close()    
+            local_file.close()    
+