Improve progress bars, soundcloud working

This commit is contained in:
Nathan Thomas 2023-11-21 16:29:31 -08:00
parent 3640e4e70a
commit f9b263a718
20 changed files with 213 additions and 86 deletions

View File

@ -3,10 +3,11 @@ import logging
import os
from dataclasses import dataclass
from . import progress
from .artwork import download_artwork
from .client import Client
from .config import Config
from .console import console
from .exceptions import NonStreamable
from .media import Media, Pending
from .metadata import AlbumMetadata
from .metadata.util import get_album_track_ids
@ -24,20 +25,19 @@ class Album(Media):
folder: str
async def preprocess(self):
if self.config.session.cli.text_output:
console.print(
f"Downloading [cyan]{self.meta.album}[/cyan] by [cyan]{self.meta.albumartist}[/cyan]"
)
progress.add_title(self.meta.album)
async def download(self):
async def _resolve_and_download(pending):
async def _resolve_and_download(pending: Pending):
track = await pending.resolve()
if track is None:
return
await track.rip()
await asyncio.gather(*[_resolve_and_download(p) for p in self.tracks])
async def postprocess(self):
pass
progress.remove_title(self.meta.album)
@dataclass(slots=True)
@ -46,9 +46,17 @@ class PendingAlbum(Pending):
client: Client
config: Config
async def resolve(self):
async def resolve(self) -> Album | None:
resp = await self.client.get_metadata(self.id, "album")
meta = AlbumMetadata.from_resp(resp, self.client.source)
try:
meta = AlbumMetadata.from_album_resp(resp, self.client.source)
except NonStreamable:
logger.error(
f"Album {self.id} not available to stream on {self.client.source}"
)
return None
tracklist = get_album_track_ids(self.client.source, resp)
folder = self.config.session.downloads.folder
album_folder = self._album_folder(folder, meta)

View File

@ -71,7 +71,7 @@ async def download_artwork(
)
_, embed_url, embed_cover_path = covers.get_size(config.embed_size)
if embed_cover_path is None and config.embed:
if embed_cover_path is None and embed:
assert embed_url is not None
embed_dir = os.path.join(folder, "__artwork")
os.makedirs(embed_dir, exist_ok=True)
@ -89,13 +89,13 @@ async def download_artwork(
await asyncio.gather(*downloadables)
# Update `covers` to reflect the current download state
if config.save_artwork:
if save_artwork:
assert saved_cover_path is not None
covers.set_largest_path(saved_cover_path)
if config.saved_max_width > 0:
downscale_image(saved_cover_path, config.saved_max_width)
if config.embed:
if embed:
assert embed_cover_path is not None
covers.set_path(config.embed_size, embed_cover_path)
if config.embed_max_width > 0:

View File

@ -153,8 +153,9 @@ class MetadataConfig:
# Sets the value of the 'ALBUM' field in the metadata to the playlist's name.
# This is useful if your music library software organizes tracks based on album name.
set_playlist_to_album: bool
# Replaces the original track's tracknumber with it's position in the playlist
new_playlist_tracknumbers: bool
# If part of a playlist, sets the `tracknumber` field in the metadata to the track's
# position in the playlist instead of its position in its album
renumber_playlist_tracks: bool
# The following metadata tags won't be applied
# See https://github.com/nathom/streamrip/wiki/Metadata-Tag-Names for more info
exclude: list[str]
@ -314,6 +315,20 @@ class ConfigData:
update_toml_section_from_config(self.toml["database"], self.database)
update_toml_section_from_config(self.toml["conversion"], self.conversion)
def get_source(
self, source: str
) -> QobuzConfig | DeezerConfig | SoundcloudConfig | TidalConfig:
d = {
"qobuz": self.qobuz,
"deezer": self.deezer,
"soundcloud": self.soundcloud,
"tidal": self.tidal,
}
res = d.get(source)
if res is None:
raise Exception(f"Invalid source {source}")
return res
def update_toml_section_from_config(toml_section, config):
for field in fields(config):

View File

@ -141,8 +141,9 @@ saved_max_width = -1
# Sets the value of the 'ALBUM' field in the metadata to the playlist's name.
# This is useful if your music library software organizes tracks based on album name.
set_playlist_to_album = true
# Replaces the original track's tracknumber with it's position in the playlist
new_playlist_tracknumbers = true
# If part of a playlist, sets the `tracknumber` field in the metadata to the track's
# position in the playlist instead of its position in its album
renumber_playlist_tracks = true
# The following metadata tags won't be applied
# See https://github.com/nathom/streamrip/wiki/Metadata-Tag-Names for more info
exclude = []

View File

@ -264,3 +264,17 @@ class AAC(Converter):
def get_quality_arg(self, _: int) -> str:
return ""
def get(codec: str) -> type[Converter]:
CONV_CLASS = {
"FLAC": FLAC,
"ALAC": ALAC,
"MP3": LAME,
"OPUS": OPUS,
"OGG": Vorbis,
"VORBIS": Vorbis,
"AAC": AAC,
"M4A": AAC,
}
return CONV_CLASS[codec.upper()]

View File

@ -42,6 +42,7 @@ class Downloadable(ABC):
async def size(self) -> int:
if self._size is not None:
return self._size
async with self.session.head(self.url) as response:
response.raise_for_status()
content_length = response.headers.get("Content-Length", 0)
@ -231,11 +232,12 @@ class SoundcloudDownloadable(Downloadable):
return tmp
async def size(self) -> int:
async with self.session.get(self.url) as resp:
content = await resp.text("utf-8")
if self.file_type == "mp3":
async with self.session.get(self.url) as resp:
content = await resp.text("utf-8")
parsed_m3u = m3u8.loads(content)
self._size = len(parsed_m3u.segments)
parsed_m3u = m3u8.loads(content)
self._size = len(parsed_m3u.segments)
return await super().size()

View File

@ -77,7 +77,6 @@ class Main:
async def rip(self):
await asyncio.gather(*[item.rip() for item in self.media])
for client in self.clients.values():
if hasattr(client, "session"):
await client.session.close()

View File

@ -27,6 +27,6 @@ class Pending(ABC):
"""A request to download a `Media` whose metadata has not been fetched."""
@abstractmethod
async def resolve(self) -> Media:
async def resolve(self) -> Media | None:
"""Fetch metadata and resolve into a downloadable `Media` object."""
raise NotImplemented

View File

@ -5,6 +5,7 @@ import re
from dataclasses import dataclass
from typing import Optional
from ..exceptions import NonStreamable
from .covers import Covers
from .util import get_quality_id, safe_get, typed
@ -114,8 +115,11 @@ class AlbumMetadata:
# Non-embedded information
# version = resp.get("version")
cover_urls = Covers.from_qobuz(resp)
streamable = typed(resp.get("streamable", False), bool)
assert streamable
# streamable = typed(resp.get("streamable", False), bool)
#
# if not streamable:
# raise NonStreamable(resp)
bit_depth = typed(resp.get("maximum_bit_depth"), int | None)
sampling_rate = typed(resp.get("maximum_sampling_rate"), int | float | None)
quality = get_quality_id(bit_depth, sampling_rate)
@ -166,7 +170,6 @@ class AlbumMetadata:
@classmethod
def from_soundcloud(cls, resp) -> AlbumMetadata:
track = resp
logger.debug(track)
track_id = track["id"]
bit_depth, sampling_rate = None, None
explicit = typed(
@ -227,7 +230,7 @@ class AlbumMetadata:
raise NotImplementedError
@classmethod
def from_resp(cls, resp: dict, source: str) -> AlbumMetadata:
def from_track_resp(cls, resp: dict, source: str) -> AlbumMetadata:
if source == "qobuz":
return cls.from_qobuz(resp["album"])
if source == "tidal":
@ -237,3 +240,15 @@ class AlbumMetadata:
if source == "deezer":
return cls.from_deezer(resp["album"])
raise Exception("Invalid source")
@classmethod
def from_album_resp(cls, resp: dict, source: str) -> AlbumMetadata:
if source == "qobuz":
return cls.from_qobuz(resp)
if source == "tidal":
return cls.from_tidal(resp)
if source == "soundcloud":
return cls.from_soundcloud(resp)
if source == "deezer":
return cls.from_deezer(resp)
raise Exception("Invalid source")

View File

@ -1,3 +1,4 @@
import logging
from dataclasses import dataclass
from .album_metadata import AlbumMetadata
@ -8,6 +9,8 @@ NON_STREAMABLE = "_non_streamable"
ORIGINAL_DOWNLOAD = "_original_download"
NOT_RESOLVED = "_not_resolved"
logger = logging.getLogger("streamrip")
def get_soundcloud_id(resp: dict) -> str:
item_id = resp["id"]
@ -44,11 +47,19 @@ class PlaylistMetadata:
@classmethod
def from_qobuz(cls, resp: dict):
name = typed(resp["title"], str)
tracks = [
TrackMetadata.from_qobuz(AlbumMetadata.from_qobuz(track["album"]), track)
for track in resp["tracks"]["items"]
]
logger.debug(resp)
name = typed(resp["name"], str)
tracks = []
for i, track in enumerate(resp["tracks"]["items"]):
meta = TrackMetadata.from_qobuz(
AlbumMetadata.from_qobuz(track["album"]), track
)
if meta is None:
logger.error(f"Track {i+1} in playlist {name} not available for stream")
continue
tracks.append(meta)
return cls(name, tracks)
@classmethod

View File

@ -3,6 +3,7 @@ from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
from ..exceptions import NonStreamable
from .album_metadata import AlbumMetadata
from .util import safe_get, typed
@ -30,8 +31,13 @@ class TrackMetadata:
composer: str | None
@classmethod
def from_qobuz(cls, album: AlbumMetadata, resp: dict) -> TrackMetadata:
def from_qobuz(cls, album: AlbumMetadata, resp: dict) -> TrackMetadata | None:
title = typed(resp["title"].strip(), str)
streamable = typed(resp.get("streamable", False), bool)
if not streamable:
return None
version = typed(resp.get("version"), str | None)
work = typed(resp.get("work"), str | None)
if version is not None and version not in title:
@ -114,7 +120,7 @@ class TrackMetadata:
raise NotImplemented
@classmethod
def from_resp(cls, album: AlbumMetadata, source, resp) -> TrackMetadata:
def from_resp(cls, album: AlbumMetadata, source, resp) -> TrackMetadata | None:
if source == "qobuz":
return cls.from_qobuz(album, resp)
if source == "tidal":

View File

@ -3,6 +3,7 @@ import logging
import os
from dataclasses import dataclass
from . import progress
from .artwork import download_artwork
from .client import Client
from .config import Config
@ -20,13 +21,25 @@ class PendingPlaylistTrack(Pending):
client: Client
config: Config
folder: str
playlist_name: str
position: int
async def resolve(self) -> Track:
async def resolve(self) -> Track | None:
resp = await self.client.get_metadata(self.id, "track")
album = AlbumMetadata.from_resp(resp["album"], self.client.source)
album = AlbumMetadata.from_resp(resp, self.client.source)
meta = TrackMetadata.from_resp(album, self.client.source, resp)
quality = getattr(self.config.session, self.client.source).quality
assert isinstance(quality, int)
if meta is None:
logger.error(f"Cannot stream track ({self.id}) on {self.client.source}")
return None
c = self.config.session.metadata
if c.renumber_playlist_tracks:
meta.tracknumber = self.position
if c.set_playlist_to_album:
album.album = self.playlist_name
quality = self.config.session.get_source(self.client.source).quality
embedded_cover_path, downloadable = await asyncio.gather(
self._download_cover(album.covers, self.folder),
self.client.get_downloadable(self.id, quality),
@ -55,11 +68,16 @@ class Playlist(Media):
pass
async def download(self):
async def _resolve_and_download(pending):
progress.add_title(self.name)
async def _resolve_and_download(pending: PendingPlaylistTrack):
track = await pending.resolve()
if track is None:
return
await track.rip()
await asyncio.gather(*[_resolve_and_download(p) for p in self.tracks])
progress.remove_title(self.name)
async def postprocess(self):
pass
@ -71,14 +89,16 @@ class PendingPlaylist(Pending):
client: Client
config: Config
async def resolve(self):
async def resolve(self) -> Playlist | None:
resp = await self.client.get_metadata(self.id, "playlist")
meta = PlaylistMetadata.from_resp(resp, self.client.source)
name = meta.name
parent = self.config.session.downloads.folder
folder = os.path.join(parent, clean_filename(name))
tracks = [
PendingPlaylistTrack(id, self.client, self.config, folder)
for id in meta.ids()
PendingPlaylistTrack(
id, self.client, self.config, folder, name, position + 1
)
for position, id in enumerate(meta.ids())
]
return Playlist(name, self.config, self.client, tracks)

View File

@ -1,54 +1,83 @@
from dataclasses import dataclass
from typing import Callable
from click import style
from rich.console import Group
from rich.live import Live
from rich.progress import Progress
from rich.text import Text
from .console import console
THEMES = {
"plain": None,
"dainty": (
"{desc} |{bar}| "
+ style("{remaining}", fg="magenta")
+ " left at "
+ style("{rate_fmt}{postfix} ", fg="cyan", bold=True)
),
}
class ProgressManager:
def __init__(self):
self.started = False
self.progress = Progress(console=console)
self.prefix = Text.assemble(("Downloading ", "bold cyan"), overflow="ellipsis")
self.live = Live(Group(self.prefix, self.progress), refresh_per_second=10)
self.task_titles = []
def get_callback(self, total: int, desc: str):
if not self.started:
self.progress.start()
self.live.start()
self.started = True
task = self.progress.add_task(f"[cyan]{desc}", total=total)
def _callback(x: int):
def _callback_update(x: int):
self.progress.update(task, advance=x)
self.live.update(Group(self.get_title_text(), self.progress))
return _callback
def _callback_done():
self.progress.update(task, visible=False)
return Handle(_callback_update, _callback_done)
def cleanup(self):
if self.started:
self.progress.stop()
self.live.stop()
def add_title(self, title: str):
self.task_titles.append(title)
def remove_title(self, title: str):
self.task_titles.remove(title)
def get_title_text(self) -> Text:
t = self.prefix + Text(", ".join(self.task_titles))
t.overflow = "ellipsis"
return t
@dataclass
class Handle:
update: Callable[[int], None]
done: Callable[[], None]
def __enter__(self):
return self.update
def __exit__(self, *_):
self.done()
# global instance
_p = ProgressManager()
def get_progress_callback(
enabled: bool, total: int, desc: str
) -> Callable[[int], None]:
def get_progress_callback(enabled: bool, total: int, desc: str) -> Handle:
if not enabled:
return lambda _: None
return Handle(lambda _: None, lambda: None)
return _p.get_callback(total, desc)
def add_title(title: str):
_p.add_title(title)
def remove_title(title: str):
_p.remove_title(title)
def clear_progress():
_p.cleanup()

View File

@ -167,9 +167,8 @@ class QobuzClient(Client):
assert status == 200
yield resp
async def get_downloadable(self, item: dict, quality: int) -> Downloadable:
async def get_downloadable(self, item_id: str, quality: int) -> Downloadable:
assert self.secret is not None and self.logged_in and 1 <= quality <= 4
item_id = item["id"]
status, resp_json = await self._request_file_url(item_id, quality, self.secret)
assert status == 200
stream_url = resp_json.get("url")

View File

@ -57,7 +57,9 @@ class SoundcloudClient(Client):
API response.
"""
if media_type == "track":
return await self._get_track(item_id)
# parse custom id that we injected
_item_id, _ = item_id.split("|")
return await self._get_track(_item_id)
elif media_type == "playlist":
return await self._get_playlist(item_id)
else:
@ -143,8 +145,10 @@ class SoundcloudClient(Client):
# if download_url == '_non_streamable' then we raise an exception
infos: list[str] = item_info.split("|")
logger.debug(f"{infos=}")
assert len(infos) == 2, infos
item_id, download_info = infos
assert re.match(r"\d+", item_id) is not None
if download_info == self.NON_STREAMABLE:
raise NonStreamable(item_info)

View File

@ -1,4 +1,5 @@
import asyncio
import logging
import os
from dataclasses import dataclass
@ -7,6 +8,7 @@ from .artwork import download_artwork
from .client import Client
from .config import Config
from .downloadable import Downloadable
from .exceptions import NonStreamable
from .filepath_utils import clean_filename
from .media import Media, Pending
from .metadata import AlbumMetadata, Covers, TrackMetadata
@ -14,6 +16,8 @@ from .progress import get_progress_callback
from .semaphore import global_download_semaphore
from .tagger import tag_file
logger = logging.getLogger("streamrip")
@dataclass(slots=True)
class Track(Media):
@ -33,12 +37,12 @@ class Track(Media):
async def download(self):
# TODO: progress bar description
async with global_download_semaphore(self.config.session.downloads):
callback = get_progress_callback(
with get_progress_callback(
self.config.session.cli.progress_bars,
await self.downloadable.size(),
f"Track {self.meta.tracknumber}",
)
await self.downloadable.download(self.download_path, callback)
) as callback:
await self.downloadable.download(self.download_path, callback)
async def postprocess(self):
await self._tag()
@ -52,19 +56,9 @@ class Track(Media):
await tag_file(self.download_path, self.meta, self.cover_path)
async def _convert(self):
CONV_CLASS: dict[str, type[converter.Converter]] = {
"FLAC": converter.FLAC,
"ALAC": converter.ALAC,
"MP3": converter.LAME,
"OPUS": converter.OPUS,
"OGG": converter.Vorbis,
"VORBIS": converter.Vorbis,
"AAC": converter.AAC,
"M4A": converter.AAC,
}
c = self.config.session.conversion
codec = c.codec
engine = CONV_CLASS[codec.upper()](
engine_class = converter.get(c.codec)
engine = engine_class(
filename=self.download_path,
sampling_rate=c.sampling_rate,
bit_depth=c.bit_depth,
@ -97,9 +91,15 @@ class PendingTrack(Pending):
# cover_path is None <==> Artwork for this track doesn't exist in API
cover_path: str | None
async def resolve(self) -> Track:
async def resolve(self) -> Track | None:
resp = await self.client.get_metadata(self.id, "track")
meta = TrackMetadata.from_resp(self.album, self.client.source, resp)
if meta is None:
logger.error(
f"Track {self.id} not available for stream on {self.client.source}"
)
return None
quality = getattr(self.config.session, self.client.source).quality
assert isinstance(quality, int)
downloadable = await self.client.get_downloadable(self.id, quality)
@ -118,13 +118,17 @@ class PendingSingle(Pending):
client: Client
config: Config
async def resolve(self) -> Track:
async def resolve(self) -> Track | None:
resp = await self.client.get_metadata(self.id, "track")
# Patch for soundcloud
# self.id = resp["id"]
album = AlbumMetadata.from_resp(resp, self.client.source)
album = AlbumMetadata.from_track_resp(resp, self.client.source)
meta = TrackMetadata.from_resp(album, self.client.source, resp)
if meta is None:
logger.error(f"Cannot stream track ({self.id}) on {self.client.source}")
return None
quality = getattr(self.config.session, self.client.source).quality
assert isinstance(quality, int)
folder = os.path.join(

View File

@ -54,6 +54,8 @@ class GenericURL(URL):
return PendingSingle(item_id, client, config)
elif media_type == "album":
return PendingAlbum(item_id, client, config)
elif media_type == "playlist":
return PendingPlaylist(item_id, client, config)
else:
raise NotImplementedError

File diff suppressed because one or more lines are too long

View File

@ -92,7 +92,7 @@ def test_sample_config_data_fields(sample_config_data):
saved_max_width=-1,
),
metadata=MetadataConfig(
set_playlist_to_album=True, new_playlist_tracknumbers=True, exclude=[]
set_playlist_to_album=True, renumber_playlist_tracks=True, exclude=[]
),
qobuz_filters=QobuzDiscographyFilterConfig(
extras=False,
@ -102,7 +102,6 @@ def test_sample_config_data_fields(sample_config_data):
non_studio_albums=False,
non_remaster=False,
),
theme=ThemeConfig(progress_bar="dainty"),
database=DatabaseConfig(
downloads_enabled=True,
downloads_path="downloadspath",
@ -130,7 +129,6 @@ def test_sample_config_data_fields(sample_config_data):
assert sample_config_data.filepaths == test_config.filepaths
assert sample_config_data.metadata == test_config.metadata
assert sample_config_data.qobuz_filters == test_config.qobuz_filters
assert sample_config_data.theme == test_config.theme
assert sample_config_data.database == test_config.database
assert sample_config_data.conversion == test_config.conversion

View File

@ -7,7 +7,7 @@ from streamrip.config import *
@pytest.fixture
def toml():
with open("streamrip/config.toml") as f:
t = tomlkit.parse(f.read())
t = tomlkit.parse(f.read()) # type: ignore
return t