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} # {url: path}
self._paths: Dict[str, str] = {} self._paths: Dict[str, str] = {}
self.task: Optional[asyncio.Task] = None self.task: Optional[asyncio.Task] = None
self.callack = chunk_callback
if tempdir is None: if tempdir is None:
tempdir = gettempdir() tempdir = gettempdir()
@ -208,7 +207,8 @@ class DownloadPool:
logger.debug("Finished %s", url) logger.debug("Finished %s", url)
def download(self): def download(self, callback=None):
self.callback = callback
asyncio.run(self._download_urls()) asyncio.run(self._download_urls())
@property @property

View file

@ -58,6 +58,7 @@ from .utils import (
get_container, get_container,
get_cover_urls, get_cover_urls,
get_stats_from_quality, get_stats_from_quality,
get_tqdm_bar,
safe_get, safe_get,
tidal_cover_url, tidal_cover_url,
tqdm_stream, tqdm_stream,
@ -459,9 +460,17 @@ class Track(Media):
with DownloadPool( with DownloadPool(
segment.uri for segment in parsed_m3u.segments segment.uri for segment in parsed_m3u.segments
) as pool: ) 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( subprocess.call(
[ (
"ffmpeg", "ffmpeg",
"-i", "-i",
f"concat:{'|'.join(pool.files)}", f"concat:{'|'.join(pool.files)}",
@ -470,7 +479,7 @@ class Track(Media):
"-loglevel", "-loglevel",
"panic", "panic",
self.path, self.path,
] )
) )
# self.path += ".mp3" # self.path += ".mp3"
@ -883,20 +892,28 @@ class Video(Media):
import m3u8 import m3u8
import requests import requests
secho( # secho(
f"Downloading {self.title} (Video). This may take a while.", # f"Downloading {self.title} (Video). This may take a while.",
fg="blue", # fg="blue",
) # )
self.parent_folder = kwargs.get("parent_folder", "StreamripDownloads") self.parent_folder = kwargs.get("parent_folder", "StreamripDownloads")
url = self.client.get_file_url(self.id, video=True) url = self.client.get_file_url(self.id, video=True)
parsed_m3u = m3u8.loads(requests.get(url).text) parsed_m3u = m3u8.loads(requests.get(url).text)
# Asynchronously download the streams # Asynchronously download the streams
with DownloadPool( with DownloadPool(
segment.uri for segment in parsed_m3u.segments segment.uri for segment in parsed_m3u.segments
) as pool: ) 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 # Put the filenames in a tempfile that ffmpeg
# can read from # can read from
@ -919,6 +936,8 @@ class Video(Media):
file_list_path, file_list_path,
"-c", "-c",
"copy", "copy",
"-loglevel",
"panic",
self.path, self.path,
] ]
) )

View file

@ -439,10 +439,10 @@ def tqdm_stream(iterator, desc: Optional[str] = None) -> Iterator[bytes]:
yield chunk yield chunk
def get_tqdm_bar(total, desc: Optional[str] = None): def get_tqdm_bar(total, desc: Optional[str] = None, unit="B"):
return tqdm( return tqdm(
total=total, total=total,
unit="B", unit=unit,
unit_scale=True, unit_scale=True,
unit_divisor=1024, unit_divisor=1024,
desc=desc, desc=desc,