From 700b435367afeb82ddee3bc325d6a3fbe9761e63 Mon Sep 17 00:00:00 2001
From: Nathan Thomas <nathanthomas707@gmail.com>
Date: Sat, 11 Sep 2021 21:31:12 -0700
Subject: [PATCH] Add tqdm bar for hls downloads

---
 streamrip/downloadtools.py |  4 ++--
 streamrip/media.py         | 35 +++++++++++++++++++++++++++--------
 streamrip/utils.py         |  4 ++--
 3 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/streamrip/downloadtools.py b/streamrip/downloadtools.py
index 310ca21..e35917a 100644
--- a/streamrip/downloadtools.py
+++ b/streamrip/downloadtools.py
@@ -172,7 +172,6 @@ class DownloadPool:
         # {url: path}
         self._paths: Dict[str, str] = {}
         self.task: Optional[asyncio.Task] = None
-        self.callack = chunk_callback
 
         if tempdir is None:
             tempdir = gettempdir()
@@ -208,7 +207,8 @@ class DownloadPool:
 
         logger.debug("Finished %s", url)
 
-    def download(self):
+    def download(self, callback=None):
+        self.callback = callback
         asyncio.run(self._download_urls())
 
     @property
diff --git a/streamrip/media.py b/streamrip/media.py
index e010dac..9b06513 100644
--- a/streamrip/media.py
+++ b/streamrip/media.py
@@ -58,6 +58,7 @@ from .utils import (
     get_container,
     get_cover_urls,
     get_stats_from_quality,
+    get_tqdm_bar,
     safe_get,
     tidal_cover_url,
     tqdm_stream,
@@ -459,9 +460,17 @@ class Track(Media):
             with DownloadPool(
                 segment.uri for segment in parsed_m3u.segments
             ) as pool:
-                pool.download()
+
+                bar = get_tqdm_bar(
+                    len(pool), desc=self._progress_desc, unit="Chunk"
+                )
+
+                def update_tqdm_bar():
+                    bar.update(1)
+
+                pool.download(callback=update_tqdm_bar)
                 subprocess.call(
-                    [
+                    (
                         "ffmpeg",
                         "-i",
                         f"concat:{'|'.join(pool.files)}",
@@ -470,7 +479,7 @@ class Track(Media):
                         "-loglevel",
                         "panic",
                         self.path,
-                    ]
+                    )
                 )
 
             # self.path += ".mp3"
@@ -883,20 +892,28 @@ class Video(Media):
         import m3u8
         import requests
 
-        secho(
-            f"Downloading {self.title} (Video). This may take a while.",
-            fg="blue",
-        )
+        # secho(
+        #     f"Downloading {self.title} (Video). This may take a while.",
+        #     fg="blue",
+        # )
 
         self.parent_folder = kwargs.get("parent_folder", "StreamripDownloads")
         url = self.client.get_file_url(self.id, video=True)
 
         parsed_m3u = m3u8.loads(requests.get(url).text)
         # Asynchronously download the streams
+
         with DownloadPool(
             segment.uri for segment in parsed_m3u.segments
         ) as pool:
-            pool.download()
+            bar = get_tqdm_bar(
+                len(pool), desc="Downloading Video", unit="Chunk"
+            )
+
+            def update_tqdm_bar():
+                bar.update(1)
+
+            pool.download(callback=update_tqdm_bar)
 
             # Put the filenames in a tempfile that ffmpeg
             # can read from
@@ -919,6 +936,8 @@ class Video(Media):
                     file_list_path,
                     "-c",
                     "copy",
+                    "-loglevel",
+                    "panic",
                     self.path,
                 ]
             )
diff --git a/streamrip/utils.py b/streamrip/utils.py
index 4d3bab7..a64b4f8 100644
--- a/streamrip/utils.py
+++ b/streamrip/utils.py
@@ -439,10 +439,10 @@ def tqdm_stream(iterator, desc: Optional[str] = None) -> Iterator[bytes]:
             yield chunk
 
 
-def get_tqdm_bar(total, desc: Optional[str] = None):
+def get_tqdm_bar(total, desc: Optional[str] = None, unit="B"):
     return tqdm(
         total=total,
-        unit="B",
+        unit=unit,
         unit_scale=True,
         unit_divisor=1024,
         desc=desc,