Optional progress bar

entsoe/__init__.py

@@ -1,2 +1,3 @@
 from .entsoe import EntsoeRawClient, EntsoePandasClient, __version__
 from .mappings import Area
+from .decorators import ProgressBar

entsoe/decorators.py

@@ -1,3 +1,4 @@
+import tqdm
 import sys
 from socket import gaierror
 from time import sleep
@@ -6,10 +7,47 @@
 from .exceptions import NoMatchingDataError, PaginationError
 import pandas as pd
 import logging
+from .misc import year_blocks, day_blocks, period_splitter
 
-from .misc import year_blocks, day_blocks
+class ProgressBar:
+    show = False
+    batch_days = 7
 
+    @staticmethod
+    def progress_bar(func):
+        """ Forces a weekly progress bar, if query range more than 7 days (adjustable interval)"""
+        @wraps(func)
+        def progress_wrapper(*args, **kwargs):
+            if not ProgressBar.show:
+                result = func(*args, **kwargs)
+                return result
+
+            start = kwargs['start']
+            end = kwargs['end']
+
+            query_bins = period_splitter(start, end, days_thresh=ProgressBar.batch_days)
+
+            if len(query_bins) <= 1:
+                result = func(*args, **kwargs)
+
+            else:
+
+                tqdm.tqdm._instances.clear()
+                pbar = tqdm.tqdm(range(len(query_bins)), desc="\tDownload Progress")
+                results = []
+                for i in pbar:
+                    pair = query_bins[i]
+                    pair_kwargs = kwargs.copy()
+                    pair_kwargs['start'] = pair[0]
+                    pair_kwargs['end'] = pair[1]
+                    results.append(func(*args, **pair_kwargs ))
+                    pbar.set_postfix_str(s="Reached: {}".format(pair[1].date()))
+
+                result = pd.concat([*results], axis = 0)
+
+            return result
 
+        return progress_wrapper
 def retry(func):
     """Catches connection errors, waits and retries"""
 

entsoe/entsoe.py

@@ -16,9 +16,8 @@
     parse_unavailabilities, parse_contracted_reserve, parse_imbalance_prices_zip, \
     parse_imbalance_volumes_zip, parse_netpositions, parse_procured_balancing_capacity, \
     parse_water_hydro
-from .decorators import retry, paginated, year_limited, day_limited, documents_limited
+from .decorators import retry, paginated, year_limited, day_limited, documents_limited, ProgressBar
 import warnings
-
 warnings.filterwarnings('ignore', category=XMLParsedAsHTMLWarning)
 
 __title__ = "entsoe-py"
@@ -306,7 +305,8 @@ def query_wind_and_solar_forecast(
         return response.text
 
     def query_intraday_wind_and_solar_forecast(
-            self, country_code: Union[Area, str], start: pd.Timestamp, end: pd.Timestamp, psr_type: Optional[str] = None) -> str:
+            self, country_code: Union[Area, str], start: pd.Timestamp, end: pd.Timestamp,
+            psr_type: Optional[str] = None) -> str:
         return self.query_wind_and_solar_forecast(country_code=country_code,
                                                   start=start,
                                                   end=end,
@@ -1032,6 +1032,7 @@ def query_withdrawn_unavailability_of_generation_units(
         return content
 
 class EntsoePandasClient(EntsoeRawClient):
+    @ProgressBar.progress_bar
     @year_limited
     def query_net_position(self, country_code: Union[Area, str],
                             start: pd.Timestamp, end: pd.Timestamp, dayahead: bool = True) -> pd.Series:
@@ -1055,6 +1056,7 @@ def query_net_position(self, country_code: Union[Area, str],
         series = series.truncate(before=start, after=end)
         return series
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_day_ahead_prices(
             self, country_code: Union[Area, str],
@@ -1093,6 +1095,7 @@ def query_day_ahead_prices(
             raise NoMatchingDataError
         return series
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_load(self, country_code: Union[Area, str], start: pd.Timestamp,
                    end: pd.Timestamp) -> pd.DataFrame:
@@ -1116,6 +1119,7 @@ def query_load(self, country_code: Union[Area, str], start: pd.Timestamp,
         df = df.truncate(before=start, after=end)
         return df
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_load_forecast(
             self, country_code: Union[Area, str], start: pd.Timestamp,
@@ -1163,6 +1167,7 @@ def query_load_and_forecast(
         return df_load_forecast_da.join(df_load, sort=True, how='inner')
 
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_generation_forecast(
             self, country_code: Union[Area, str], start: pd.Timestamp,
@@ -1190,6 +1195,7 @@ def query_generation_forecast(
         df = df.truncate(before=start, after=end)
         return df
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_wind_and_solar_forecast(
             self, country_code: Union[Area, str], start: pd.Timestamp,
@@ -1218,6 +1224,7 @@ def query_wind_and_solar_forecast(
         df = df.truncate(before=start, after=end)
         return df
 
+    @ProgressBar.progress_bar
     def query_intraday_wind_and_solar_forecast(
             self, country_code: Union[Area, str], start: pd.Timestamp,
             end: pd.Timestamp, psr_type: Optional[str] = None) -> pd.DataFrame:
@@ -1228,6 +1235,8 @@ def query_intraday_wind_and_solar_forecast(
                                                   process_type='A40')
 
 
+
+    @ProgressBar.progress_bar
     @year_limited
     def query_generation(
             self, country_code: Union[Area, str], start: pd.Timestamp,
@@ -1256,6 +1265,7 @@ def query_generation(
         df = df.truncate(before=start, after=end)
         return df
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_installed_generation_capacity(
             self, country_code: Union[Area, str], start: pd.Timestamp,
@@ -1283,6 +1293,7 @@ def query_installed_generation_capacity(
         df = df.truncate(before=start - YearBegin(), after=end + YearEnd())
         return df
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_installed_generation_capacity_per_unit(
             self, country_code: Union[Area, str], start: pd.Timestamp,
@@ -1323,6 +1334,7 @@ def query_aggregate_water_reservoirs_and_hydro_storage(self, country_code: Union
         return df
 
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_crossborder_flows(
             self, country_code_from: Union[Area, str],
@@ -1354,6 +1366,7 @@ def query_crossborder_flows(
         ts = ts.truncate(before=start, after=end)
         return ts
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_scheduled_exchanges(
             self, country_code_from: Union[Area, str],
@@ -1390,6 +1403,7 @@ def query_scheduled_exchanges(
         ts = ts.truncate(before=start, after=end)
         return ts
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_net_transfer_capacity_dayahead(
             self, country_code_from: Union[Area, str],
@@ -1421,6 +1435,7 @@ def query_net_transfer_capacity_dayahead(
         ts = ts.truncate(before=start, after=end)
         return ts
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_net_transfer_capacity_weekahead(
             self, country_code_from: Union[Area, str],
@@ -1452,6 +1467,7 @@ def query_net_transfer_capacity_weekahead(
         ts = ts.truncate(before=start, after=end)
         return ts
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_net_transfer_capacity_monthahead(
             self, country_code_from: Union[Area, str],
@@ -1483,6 +1499,7 @@ def query_net_transfer_capacity_monthahead(
         ts = ts.truncate(before=start, after=end)
         return ts
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_net_transfer_capacity_yearahead(
             self, country_code_from: Union[Area, str],
@@ -1514,6 +1531,7 @@ def query_net_transfer_capacity_yearahead(
         ts = ts.truncate(before=start, after=end)
         return ts
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_intraday_offered_capacity(
         self, country_code_from: Union[Area, str],
@@ -1591,6 +1609,7 @@ def query_offered_capacity(
         ts = ts.truncate(before=start, after=end)
         return ts
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_imbalance_prices(
             self, country_code: Union[Area, str], start: pd.Timestamp,
@@ -1616,6 +1635,7 @@ def query_imbalance_prices(
         df = df.truncate(before=start, after=end)
         return df
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_imbalance_volumes(
             self, country_code: Union[Area, str], start: pd.Timestamp,
@@ -1675,6 +1695,7 @@ def query_procured_balancing_capacity(
         df = df.truncate(before=start, after=end)
         return df
 
+    @ProgressBar.progress_bar
     @year_limited
     def query_activated_balancing_energy(
             self, country_code: Union[Area, str], start: pd.Timestamp,
@@ -1925,6 +1946,7 @@ def query_withdrawn_unavailability_of_generation_units(
         df = df[(df['start'] < end) | (df['end'] > start)]
         return df
 
+    @ProgressBar.progress_bar
     @day_limited
     def query_generation_per_plant(
             self, country_code: Union[Area, str], start: pd.Timestamp,
@@ -1963,6 +1985,8 @@ def query_generation_per_plant(
         df = df.truncate(before=start, after=end)
         return df
 
+
+    @ProgressBar.progress_bar
     def query_physical_crossborder_allborders(self, country_code: Union[Area, str], start: pd.Timestamp,
                      end: pd.Timestamp, export: bool, per_hour: bool = False) -> pd.DataFrame:
         """
@@ -2001,9 +2025,9 @@ def query_physical_crossborder_allborders(self, country_code: Union[Area, str],
         df['sum'] = df.sum(axis=1)
         if per_hour:
             df = df.resample('h').mean()
-
         return df
 
+    @ProgressBar.progress_bar
     def query_import(self, country_code: Union[Area, str], start: pd.Timestamp,
                      end: pd.Timestamp) -> pd.DataFrame:
         """
@@ -2013,7 +2037,7 @@ def query_import(self, country_code: Union[Area, str], start: pd.Timestamp,
                                                    start=start,
                                                    end=end,
                                                    export=False)
-
+    @ProgressBar.progress_bar
     def query_generation_import(
             self, country_code: Union[Area, str], start: pd.Timestamp,
             end: pd.Timestamp) -> pd.DataFrame:

entsoe/misc.py

@@ -2,6 +2,23 @@
 from dateutil import rrule
 from itertools import tee
 
+def period_splitter(start, end, days_thresh = 7):
+    days_requested = (end - start).days
+    request_bins = []
+    rolling_start = start
+    if days_requested > days_thresh:
+        rolling_end = rolling_start + pd.Timedelta(days_thresh, 'D')
+        while rolling_end < end:
+            request_bins.append((rolling_start, rolling_end))
+            rolling_start += pd.Timedelta(days_thresh, 'D')
+            rolling_end += pd.Timedelta(days_thresh, 'D')
+
+        request_bins.append((rolling_start, end))
+    else:
+        request_bins = [(start, end)]
+
+    return request_bins
+
 
 def year_blocks(start, end):
     """

requirements.txt

@@ -2,3 +2,4 @@ requests
 pytz
 beautifulsoup4>=4.11.1
 pandas>=1.4.0
+tqdm>=4.65.0

setup.py

@@ -67,7 +67,7 @@
 
     # List run-time dependencies here.  These will be installed by pip when
     # your project is installed.
-    install_requires=['requests', 'pytz', 'beautifulsoup4>=4.11.1', 'pandas>=1.4.0'],
+    install_requires=['requests', 'pytz', 'beautifulsoup4>=4.11.1', 'pandas>=1.4.0', 'tqdm'],
 
     include_package_data=True,
 )

test.py

@@ -1,11 +1,16 @@
+import time
+
 import pandas as pd
 from settings import api_key
-from entsoe import EntsoePandasClient as Entsoe
+from entsoe import EntsoePandasClient as Entsoe, ProgressBar
+
+ProgressBar.show = True
+ProgressBar.batch_days = 30
+
 e = Entsoe(api_key=api_key, retry_count=20, retry_delay=30)
 
 start = pd.Timestamp('20170601', tz='Europe/Brussels')
 end = pd.Timestamp('20171201', tz='Europe/Brussels')
-
 #s = e.query_imbalance_prices(country_code='BE', start=start, end=end, as_dataframe=True)
 
 """domains = [["10YIT-GRTN-----B", "Italy, IT CA / MBA"],