Added progress bar for downloading data

pyatmos/msise/nrlmsise00.py

@@ -1039,12 +1039,12 @@ def nrlmsise00(t,lat,lon,alt,SW_OBS_PRE,omode='NoOxygen',aphmode='NoAph'):
               'f107A':f107A,'f107':f107,'ap':ap,'ap_a':aph}
     
     switches = np.ones(23)
-    if aphmode is 'Aph':
+    if aphmode == 'Aph':
         switches[8] = -1 # -1 indicates the use of 3h geomagnetic index
         
-    if omode is 'Oxygen':
+    if omode == 'Oxygen':
         output = gtd7d(inputp,switches)
-    elif omode is 'NoOxygen':
+    elif omode == 'NoOxygen':
         output = gtd7(inputp,switches)
     else:
         raise Exception("'{}' should be either 'Oxygen' or 'NoOxygen'".format(o))

pyatmos/msise/spaceweather.py

@@ -19,6 +19,8 @@ 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  =================== #
 
@@ -56,19 +58,21 @@ 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):
-        print('Downloading the latest space weather data from celestrak',end=' ... ')
+        desc = 'Downloading the latest space weather data from CELESTRAK'
-        urlretrieve(url, swfile)
+        with TqdmUpTo(unit='B', unit_scale=True, desc=desc,bar_format = bar_format) as t:
-        print('Finished')
+            urlretrieve(url, swfile,reporthook=t.update_to)
+            t.total = t.n
 
     else:
         modified_time = datetime.fromtimestamp(path.getmtime(swfile))
         if datetime.now() > modified_time + timedelta(days=1):
             remove(swfile)
-            print('Updating the space weather data from celestrak',end=' ... ')
+            desc = 'Updating the space weather data from CELESTRAK'
-            urlretrieve(url, swfile)    
+            with TqdmUpTo(unit='B', unit_scale=True, desc=desc,bar_format = bar_format) as t:
-            print('Finished')
+                urlretrieve(url, swfile,reporthook=t.update_to)    
+                t.total = t.n
         else:
             print('The space weather data in {:s} is already the latest.'.format(direc))   
     return swfile

pyatmos/msise/tqdmupto.py

@@ -0,0 +1,19 @@
+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

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup
 
 setup(
     name='pyatmos',
-    version='1.1.0',
+    version='1.1.1',
     long_description_content_type='text/markdown',
     description='A package to estimate the atmosphere parameters',
     long_description=open('README.md', 'rb').read().decode('utf-8'),
@@ -14,7 +14,7 @@ setup(
     classifiers=[
         'Intended Audience :: Education',
         'Intended Audience :: Science/Research',
-        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3.8',
         'License :: OSI Approved :: MIT License',
         ],
     packages=setuptools.find_packages(),
@@ -26,6 +26,8 @@ setup(
         'scipy',
         'numpy',
         'pyshtools',
-        'astropy'
+        'astropy',
+        'tqdm',
+        'colorama'
         ],
 )