From 033de21ed9eddd7b576d872918397e723bc22ec6 Mon Sep 17 00:00:00 2001
From: Robert Grzesik <rob.grzesik@gmail.com>
Date: Mon, 18 Dec 2023 22:05:01 -0500
Subject: [PATCH] added progress bar when downloading data from polygon

---
 lumibot/tools/polygon_helper.py | 20 ++++++++++++++++----
 requirements.txt                |  1 +
 setup.py                        |  3 ++-
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/lumibot/tools/polygon_helper.py b/lumibot/tools/polygon_helper.py
index 9d7e0efa7..8931beac0 100644
--- a/lumibot/tools/polygon_helper.py
+++ b/lumibot/tools/polygon_helper.py
@@ -6,12 +6,12 @@
 
 import pandas as pd
 import pandas_market_calendars as mcal
+from lumibot import LUMIBOT_CACHE_FOLDER
+from lumibot.entities import Asset
 
 # noinspection PyPackageRequirements
 from polygon import RESTClient
-
-from lumibot import LUMIBOT_CACHE_FOLDER
-from lumibot.entities import Asset
+from tqdm import tqdm
 
 WAIT_TIME = 60
 POLYGON_QUERY_COUNT = 0  # This is a variable that updates every time we query Polygon
@@ -73,7 +73,7 @@ def get_price_data_from_polygon(
     if not missing_dates:
         return df_all
 
-    print(f"\nGetting pricing data for {asset} / {quote_asset} with '{timespan}' timespan from Polygon...")
+    # print(f"\nGetting pricing data for {asset} / {quote_asset} with '{timespan}' timespan from Polygon...")
 
     # RESTClient connection for Polygon Stock-Equity API; traded_asset is standard
     # Add "trace=True" to see the API calls printed to the console for debugging
@@ -86,6 +86,12 @@ def get_price_data_from_polygon(
     poly_start = missing_dates[0]  # Data will start at 8am UTC (4am EST)
     poly_end = missing_dates[-1]  # Data will end at 23:59 UTC (7:59pm EST)
 
+    # Initialize tqdm progress bar
+    total_days = (missing_dates[-1] - missing_dates[0]).days + 1
+    total_queries = (total_days // MAX_POLYGON_DAYS) + 1
+    description = f"\nFetching data for {asset} / {quote_asset} '{timespan}' from Polygon..."
+    pbar = tqdm(total=total_queries, desc=description, dynamic_ncols=True)
+
     # Polygon only returns 50k results per query (~30days of 24hr 1min-candles) so we need to break up the query into
     # multiple queries if we are requesting more than 30 days of data
     delta = timedelta(days=MAX_POLYGON_DAYS)
@@ -119,12 +125,18 @@ def get_price_data_from_polygon(
             limit=50000,  # Max limit for Polygon
         )
 
+        # Update progress bar after each query
+        pbar.update(1)
+
         if result:
             df_all = update_polygon_data(df_all, result)
 
         poly_start = poly_end + timedelta(days=1)
         poly_end = poly_start + delta
 
+    # Close the progress bar when done
+    pbar.close()
+
     update_cache(cache_file, df_all, df_feather)
     return df_all
 
diff --git a/requirements.txt b/requirements.txt
index 0d88bf6c5..e179218ed 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -28,3 +28,4 @@ apscheduler
 alpaca-py
 appdirs
 pyarrow
+tqdm
diff --git a/setup.py b/setup.py
index 43179d544..bff9c5f55 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="lumibot",
-    version="2.9.9",
+    version="2.9.10",
     author="Robert Grzesik",
     author_email="rob@lumiwealth.com",
     description="Backtesting and Trading Library, Made by Lumiwealth",
@@ -44,6 +44,7 @@
         "apscheduler==3.10.4",
         "appdirs",
         "pyarrow",
+        "tqdm",
     ],
     classifiers=[
         "Programming Language :: Python :: 3",