From ae430dbd2523ce861e11b5a4e0f54d2a6dbb1fa0 Mon Sep 17 00:00:00 2001
From: Chunxiao Li <lcx366@126.com>
Date: Wed, 22 Jul 2020 09:18:58 +0800
Subject: [PATCH] Added progress bar for downloading data

---
 .DS_Store                     | Bin 6148 -> 6148 bytes
 pyatmos/.DS_Store             | Bin 8196 -> 8196 bytes
 pyatmos/msise/nrlmsise00.py   |   6 +++---
 pyatmos/msise/spaceweather.py |  18 +++++++++++-------
 pyatmos/msise/tqdmupto.py     |  19 +++++++++++++++++++
 setup.py                      |   8 +++++---
 6 files changed, 38 insertions(+), 13 deletions(-)
 create mode 100644 pyatmos/msise/tqdmupto.py

diff --git a/.DS_Store b/.DS_Store
index 9722ba2ab97500f54fab6ee4a34ae075f7af275a..abc5ca3fc6f01304389b144f2618c84b9137f9b6 100644
GIT binary patch
delta 19
acmZoMXffDuhl$P7Oh>`cWb-2?B~buC%?0ZK

delta 19
acmZoMXffDuhl$O?NJqiYV)G*=B~buCvIXh@

diff --git a/pyatmos/.DS_Store b/pyatmos/.DS_Store
index 3b33ce0c687434ece7583e71e0f8989ec58b68c2..6ac62f06c0be7892a0d4d289626f33569b0cd722 100644
GIT binary patch
delta 41
lcmZp1XmQx!D#UJSrlVkJVm5ibfaB&~A%8{`0Tq!4+yME%3+Mm<

delta 41
lcmZp1XmQx!D#UJKq@!SHVK#ZafaB&~A%8{`0Tq!4+yMD*3+4a-

diff --git a/pyatmos/msise/nrlmsise00.py b/pyatmos/msise/nrlmsise00.py
index 61bddbf..4d2e1ed 100644
--- a/pyatmos/msise/nrlmsise00.py
+++ b/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))
diff --git a/pyatmos/msise/spaceweather.py b/pyatmos/msise/spaceweather.py
index 6e6207f..7bcfa00 100644
--- a/pyatmos/msise/spaceweather.py
+++ b/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=' ... ')
-        urlretrieve(url, swfile)
-        print('Finished')
+        desc = 'Downloading the latest space weather data from CELESTRAK'
+        with TqdmUpTo(unit='B', unit_scale=True, desc=desc,bar_format = bar_format) as t:
+            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=' ... ')
-            urlretrieve(url, swfile)    
-            print('Finished')
+            desc = 'Updating the space weather data from CELESTRAK'
+            with TqdmUpTo(unit='B', unit_scale=True, desc=desc,bar_format = bar_format) as t:
+                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
diff --git a/pyatmos/msise/tqdmupto.py b/pyatmos/msise/tqdmupto.py
new file mode 100644
index 0000000..5a7ba59
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 6525a04..df7cfb4 100755
--- a/setup.py
+++ b/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'
         ],
 )