From 692ac728603c25ada6580bef63418b7b464c95ea Mon Sep 17 00:00:00 2001 From: Mandera Date: Sun, 12 Jun 2022 19:33:01 +0800 Subject: [PATCH 1/5] Initial version for progress bar when downloading --- Pipfile | 1 + webdriver_manager/core/http.py | 15 ++++++++++++++- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/Pipfile b/Pipfile index 2654d893..5a41e9b5 100644 --- a/Pipfile +++ b/Pipfile @@ -16,3 +16,4 @@ pytest-xdist = "*" [packages] requests = "*" python-dotenv = "*" +tqdm = "*" diff --git a/webdriver_manager/core/http.py b/webdriver_manager/core/http.py index 6c4922f1..f3b2be89 100644 --- a/webdriver_manager/core/http.py +++ b/webdriver_manager/core/http.py @@ -1,6 +1,8 @@ import requests from requests import Response +from tqdm import tqdm + from webdriver_manager.core.config import ssl_verify @@ -29,6 +31,17 @@ def __init__(self): self._ssl_verify = ssl_verify() def get(self, url, **kwargs) -> Response: - resp = requests.get(url=url, verify=self._ssl_verify, **kwargs) + resp = requests.get(url=url, verify=self._ssl_verify, stream=True, **kwargs) self.validate_response(resp) + + total = int(resp.headers['Content-Length']) + if total > 100: + content = bytearray() + pbar = tqdm(total=total) + for chunk in resp.iter_content(chunk_size=8192): + if chunk: # filter out keep-alive new chunks + pbar.update(len(chunk)) + content.extend(chunk) + resp._content = content # To allow content to be "consumed" again + return resp From 9e50c9f486ddb4f8c3214886e975dd3a1a97c2c7 Mon Sep 17 00:00:00 2001 From: Mandera Date: Sun, 12 Jun 2022 19:38:27 +0800 Subject: [PATCH 2/5] Trying to trigger test actions --- webdriver_manager/core/http.py | 1 + 1 file changed, 1 insertion(+) diff --git a/webdriver_manager/core/http.py b/webdriver_manager/core/http.py index f3b2be89..3807c3ec 100644 --- a/webdriver_manager/core/http.py +++ b/webdriver_manager/core/http.py @@ -45,3 +45,4 @@ def get(self, url, **kwargs) -> Response: resp._content = content # To allow content to be "consumed" again return resp + From a1dce615e11ba0c98b8e54fbad602f176d829e36 Mon Sep 17 00:00:00 2001 From: Mandera Date: Sun, 12 Jun 2022 19:56:19 +0800 Subject: [PATCH 3/5] Only display download bar if content-length is present in header --- webdriver_manager/core/http.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/webdriver_manager/core/http.py b/webdriver_manager/core/http.py index 3807c3ec..e2847c70 100644 --- a/webdriver_manager/core/http.py +++ b/webdriver_manager/core/http.py @@ -34,7 +34,7 @@ def get(self, url, **kwargs) -> Response: resp = requests.get(url=url, verify=self._ssl_verify, stream=True, **kwargs) self.validate_response(resp) - total = int(resp.headers['Content-Length']) + total = int(resp.headers.get("Content-Length", 0)) if total > 100: content = bytearray() pbar = tqdm(total=total) From cd4dfbbe60be6e487d53a86d3c91731ae9f41890 Mon Sep 17 00:00:00 2001 From: Mandera Date: Sun, 12 Jun 2022 21:20:18 +0800 Subject: [PATCH 4/5] Changed unit to bytes and automatically scale --- webdriver_manager/core/http.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/webdriver_manager/core/http.py b/webdriver_manager/core/http.py index e2847c70..6079d8fa 100644 --- a/webdriver_manager/core/http.py +++ b/webdriver_manager/core/http.py @@ -37,10 +37,10 @@ def get(self, url, **kwargs) -> Response: total = int(resp.headers.get("Content-Length", 0)) if total > 100: content = bytearray() - pbar = tqdm(total=total) + progress_bar = tqdm(total=total, unit_scale=True, unit_divisor=1024, unit="B") for chunk in resp.iter_content(chunk_size=8192): if chunk: # filter out keep-alive new chunks - pbar.update(len(chunk)) + progress_bar.update(len(chunk)) content.extend(chunk) resp._content = content # To allow content to be "consumed" again From 4cdec9a2610cef240ca3820b25cc72357786677d Mon Sep 17 00:00:00 2001 From: Mandera Date: Sun, 26 Jun 2022 12:17:51 +0800 Subject: [PATCH 5/5] Moved progress bar code to show_download_progress() in utils.py. Hardcoded the "[WDM] - " prefix for the progress bar as I think making tqdm use the logger's formatting dynamically isn't feasible. Added tqdm to install_requires in setup.py (In addition to Pipfile). --- setup.py | 3 ++- webdriver_manager/core/http.py | 15 ++------------- webdriver_manager/core/utils.py | 16 ++++++++++++++++ 3 files changed, 20 insertions(+), 14 deletions(-) diff --git a/setup.py b/setup.py index d05ae4ba..5c215db6 100644 --- a/setup.py +++ b/setup.py @@ -47,6 +47,7 @@ ], install_requires=[ 'requests', - 'python-dotenv' + 'python-dotenv', + 'tqdm' ], ) diff --git a/webdriver_manager/core/http.py b/webdriver_manager/core/http.py index 6079d8fa..bcd6cb71 100644 --- a/webdriver_manager/core/http.py +++ b/webdriver_manager/core/http.py @@ -1,9 +1,8 @@ import requests from requests import Response -from tqdm import tqdm - from webdriver_manager.core.config import ssl_verify +from webdriver_manager.core.utils import show_download_progress class HttpClient: @@ -33,16 +32,6 @@ def __init__(self): def get(self, url, **kwargs) -> Response: resp = requests.get(url=url, verify=self._ssl_verify, stream=True, **kwargs) self.validate_response(resp) - - total = int(resp.headers.get("Content-Length", 0)) - if total > 100: - content = bytearray() - progress_bar = tqdm(total=total, unit_scale=True, unit_divisor=1024, unit="B") - for chunk in resp.iter_content(chunk_size=8192): - if chunk: # filter out keep-alive new chunks - progress_bar.update(len(chunk)) - content.extend(chunk) - resp._content = content # To allow content to be "consumed" again - + show_download_progress(resp) return resp diff --git a/webdriver_manager/core/utils.py b/webdriver_manager/core/utils.py index 1df0efd3..acecc8fa 100644 --- a/webdriver_manager/core/utils.py +++ b/webdriver_manager/core/utils.py @@ -5,6 +5,8 @@ import subprocess import sys +from tqdm import tqdm + from webdriver_manager.core.archive import Archive from webdriver_manager.core.logger import log @@ -256,3 +258,17 @@ def determine_powershell(): ) as stream: stdout = stream.communicate()[0].decode() return "" if stdout == "powershell" else "powershell" + + +def show_download_progress(response, _bytes_threshold=100): + """ Opens up a response's content in chunks to show a progress bar with tqdm. + Resets response._content when done so that response can be consumed again as normal. """ + total = int(response.headers.get("Content-Length", 0)) + if total > _bytes_threshold: + content = bytearray() + progress_bar = tqdm(desc="[WDM] - Downloading", total=total, unit_scale=True, unit_divisor=1024, unit="B") + for chunk in response.iter_content(chunk_size=8192): + if chunk: # Filter out keep-alive new chunks + progress_bar.update(len(chunk)) + content.extend(chunk) + response._content = content # To allow content to be "consumed" again