Add tqdm bar for hls downloads

This commit is contained in:
Nathan Thomas 2021-09-11 21:31:12 -07:00
parent 9f5cd49aab
commit 700b435367
3 changed files with 31 additions and 12 deletions

View file

@ -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

View file

@ -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,
]
)

View file

@ -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,