Skip to content

Commit

Permalink
More robust error handling (#678)
Browse files Browse the repository at this point in the history
* Handle no copyright case for tidal

* Add default values for get calls

* Fix LSP errors

* Misc fixes
  • Loading branch information
nathom committed May 12, 2024
1 parent 868a8ff commit 527b52c
Show file tree
Hide file tree
Showing 12 changed files with 175 additions and 72 deletions.
16 changes: 7 additions & 9 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,10 @@ repository = "https://github.com/nathom/streamrip"
include = ["src/config.toml"]
keywords = ["hi-res", "free", "music", "download"]
classifiers = [
"License :: OSI Approved :: GNU General Public License (GPL)",
"Operating System :: OS Independent",
]
packages = [
{ include = "streamrip" }
"License :: OSI Approved :: GNU General Public License (GPL)",
"Operating System :: OS Independent",
]
packages = [{ include = "streamrip" }]

[tool.poetry.scripts]
rip = "streamrip.rip:rip"
Expand All @@ -25,9 +23,9 @@ python = ">=3.10 <4.0"
mutagen = "^1.45.1"
tomlkit = "^0.7.2"
pathvalidate = "^2.4.1"
simple-term-menu = {version = "^1.2.1", platform = 'darwin|linux'}
pick = {version = "^2", platform = 'win32|cygwin'}
windows-curses = {version = "^2.2.0", platform = 'win32|cygwin'}
simple-term-menu = { version = "^1.2.1", platform = 'darwin|linux' }
pick = { version = "^2", platform = 'win32|cygwin' }
windows-curses = { version = "^2.2.0", platform = 'win32|cygwin' }
Pillow = ">=9,<11"
deezer-py = "1.3.6"
pycryptodomex = "^3.10.1"
Expand Down Expand Up @@ -58,7 +56,7 @@ pytest = "^7.4"
[tool.pytest.ini_options]
minversion = "6.0"
addopts = "-ra -q"
testpaths = [ "tests" ]
testpaths = ["tests"]
log_level = "DEBUG"
asyncio_mode = 'auto'
log_cli = true
Expand Down
82 changes: 57 additions & 25 deletions streamrip/client/downloadable.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,13 +37,38 @@ def generate_temp_path(url: str):
)


async def fast_async_download(path, url, headers, callback):
"""Synchronous download with yield for every 1MB read.
Using aiofiles/aiohttp resulted in a yield to the event loop for every 1KB,
which made file downloads CPU-bound. This resulted in a ~10MB max total download
speed. This fixes the issue by only yielding to the event loop for every 1MB read.
"""
chunk_size: int = 2**17 # 131 KB
counter = 0
yield_every = 8 # 1 MB
with open(path, "wb") as file: # noqa: ASYNC101
with requests.get( # noqa: ASYNC100
url,
headers=headers,
allow_redirects=True,
stream=True,
) as resp:
for chunk in resp.iter_content(chunk_size=chunk_size):
file.write(chunk)
callback(len(chunk))
if counter % yield_every == 0:
await asyncio.sleep(0)
counter += 1


@dataclass(slots=True)
class Downloadable(ABC):
session: aiohttp.ClientSession
url: str
extension: str
chunk_size = 2**17
_size: Optional[int] = None
source: str = "Unknown"
_size_base: Optional[int] = None

async def download(self, path: str, callback: Callable[[int], Any]):
await self._download(path, callback)
Expand All @@ -58,6 +83,14 @@ async def size(self) -> int:
self._size = int(content_length)
return self._size

@property
def _size(self):
return self._size_base

@_size.setter
def _size(self, v):
self._size_base = v

@abstractmethod
async def _download(self, path: str, callback: Callable[[int], None]):
raise NotImplementedError
Expand All @@ -66,35 +99,31 @@ async def _download(self, path: str, callback: Callable[[int], None]):
class BasicDownloadable(Downloadable):
"""Just downloads a URL."""

def __init__(self, session: aiohttp.ClientSession, url: str, extension: str):
def __init__(
self,
session: aiohttp.ClientSession,
url: str,
extension: str,
source: str | None = None,
):
self.session = session
self.url = url
self.extension = extension
self._size = None
self.source: str = source or "Unknown"

async def _download(self, path: str, callback):
# Attempt to fix async performance issues by manually and infrequently
# yielding to event loop selector
counter = 0
yield_every = 16
with open(path, "wb") as file:
with requests.get(self.url, allow_redirects=True, stream=True) as resp:
for chunk in resp.iter_content(chunk_size=self.chunk_size):
file.write(chunk)
callback(len(chunk))
if counter % yield_every == 0:
await asyncio.sleep(0)
counter += 1
await fast_async_download(path, self.url, self.session.headers, callback)


class DeezerDownloadable(Downloadable):
is_encrypted = re.compile("/m(?:obile|edia)/")
# chunk_size = 2048 * 3

def __init__(self, session: aiohttp.ClientSession, info: dict):
logger.debug("Deezer info for downloadable: %s", info)
self.session = session
self.url = info["url"]
self.source: str = "deezer"
max_quality_available = max(
i for i, size in enumerate(info["quality_to_size"]) if size > 0
)
Expand Down Expand Up @@ -125,11 +154,9 @@ async def _download(self, path: str, callback):

if self.is_encrypted.search(self.url) is None:
logger.debug(f"Deezer file at {self.url} not encrypted.")
async with aiofiles.open(path, "wb") as file:
async for chunk in resp.content.iter_chunked(self.chunk_size):
await file.write(chunk)
# typically a bar.update()
callback(len(chunk))
await fast_async_download(
path, self.url, self.session.headers, callback
)
else:
blowfish_key = self._generate_blowfish_key(self.id)
logger.debug(
Expand All @@ -144,10 +171,11 @@ async def _download(self, path: str, callback):
buf += data
callback(len(data))

encrypt_chunk_size = 3 * 2048
async with aiofiles.open(path, "wb") as audio:
buflen = len(buf)
for i in range(0, buflen, self.chunk_size):
data = buf[i : min(i + self.chunk_size, buflen)]
for i in range(0, buflen, encrypt_chunk_size):
data = buf[i : min(i + encrypt_chunk_size, buflen)]
if len(data) >= 2048:
decrypted_chunk = (
self._decrypt_chunk(blowfish_key, data[:2048])
Expand Down Expand Up @@ -199,6 +227,7 @@ def __init__(
restrictions,
):
self.session = session
self.source = "tidal"
codec = codec.lower()
if codec in ("flac", "mqa"):
self.extension = "flac"
Expand All @@ -217,7 +246,7 @@ def __init__(
)
self.url = url
self.enc_key = encryption_key
self.downloadable = BasicDownloadable(session, url, self.extension)
self.downloadable = BasicDownloadable(session, url, self.extension, "tidal")

async def _download(self, path: str, callback):
await self.downloadable._download(path, callback)
Expand Down Expand Up @@ -276,6 +305,7 @@ class SoundcloudDownloadable(Downloadable):
def __init__(self, session, info: dict):
self.session = session
self.file_type = info["type"]
self.source = "soundcloud"
if self.file_type == "mp3":
self.extension = "mp3"
elif self.file_type == "original":
Expand All @@ -291,7 +321,9 @@ async def _download(self, path, callback):
await self._download_original(path, callback)

async def _download_original(self, path: str, callback):
downloader = BasicDownloadable(self.session, self.url, "flac")
downloader = BasicDownloadable(
self.session, self.url, "flac", source="soundcloud"
)
await downloader.download(path, callback)
self.size = downloader.size
engine = converter.FLAC(path)
Expand Down
14 changes: 6 additions & 8 deletions streamrip/client/qobuz.py
Original file line number Diff line number Diff line change
Expand Up @@ -197,14 +197,14 @@ async def login(self):

self.logged_in = True

async def get_metadata(self, item_id: str, media_type: str):
async def get_metadata(self, item: str, media_type: str):
if media_type == "label":
return await self.get_label(item_id)
return await self.get_label(item)

c = self.config.session.qobuz
params = {
"app_id": c.app_id,
f"{media_type}_id": item_id,
f"{media_type}_id": item,
# Do these matter?
"limit": 500,
"offset": 0,
Expand Down Expand Up @@ -302,9 +302,9 @@ async def get_user_playlists(self, limit: int = 500) -> list[dict]:
epoint = "playlist/getUserPlaylists"
return await self._paginate(epoint, {}, limit=limit)

async def get_downloadable(self, item_id: str, quality: int) -> Downloadable:
async def get_downloadable(self, item: str, quality: int) -> Downloadable:
assert self.secret is not None and self.logged_in and 1 <= quality <= 4
status, resp_json = await self._request_file_url(item_id, quality, self.secret)
status, resp_json = await self._request_file_url(item, quality, self.secret)
assert status == 200
stream_url = resp_json.get("url")

Expand All @@ -319,9 +319,7 @@ async def get_downloadable(self, item_id: str, quality: int) -> Downloadable:
raise NonStreamableError

return BasicDownloadable(
self.session,
stream_url,
"flac" if quality > 1 else "mp3",
self.session, stream_url, "flac" if quality > 1 else "mp3", source="qobuz"
)

async def _paginate(
Expand Down
7 changes: 6 additions & 1 deletion streamrip/media/album.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,12 @@ async def resolve(self) -> Album | None:
)
return None

meta = AlbumMetadata.from_album_resp(resp, self.client.source)
try:
meta = AlbumMetadata.from_album_resp(resp, self.client.source)
except Exception as e:
logger.error(f"Error building album metadata for {id=}: {e}")
return None

if meta is None:
logger.error(
f"Album {self.id} not available to stream on {self.client.source}",
Expand Down
9 changes: 8 additions & 1 deletion streamrip/media/artist.py
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,14 @@ async def resolve(self) -> Artist | None:
)
return None

meta = ArtistMetadata.from_resp(resp, self.client.source)
try:
meta = ArtistMetadata.from_resp(resp, self.client.source)
except Exception as e:
logger.error(
f"Error building artist metadata: {e}",
)
return None

albums = [
PendingAlbum(album_id, self.client, self.config, self.db)
for album_id in meta.album_ids()
Expand Down
6 changes: 5 additions & 1 deletion streamrip/media/artwork.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,11 @@ async def download_artwork(
if len(downloadables) == 0:
return embed_cover_path, saved_cover_path

await asyncio.gather(*downloadables)
try:
await asyncio.gather(*downloadables)
except Exception as e:
logger.error(f"Error downloading artwork: {e}")
return None, None

# Update `covers` to reflect the current download state
if save_artwork:
Expand Down
19 changes: 16 additions & 3 deletions streamrip/media/label.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,18 @@
import asyncio
import logging
from dataclasses import dataclass

from streamrip.exceptions import NonStreamableError

from ..client import Client
from ..config import Config
from ..db import Database
from ..metadata import LabelMetadata
from .album import PendingAlbum
from .media import Media, Pending

logger = logging.getLogger("streamrip")


@dataclass(slots=True)
class Label(Media):
Expand Down Expand Up @@ -57,9 +62,17 @@ class PendingLabel(Pending):
config: Config
db: Database

async def resolve(self) -> Label:
resp = await self.client.get_metadata(self.id, "label")
meta = LabelMetadata.from_resp(resp, self.client.source)
async def resolve(self) -> Label | None:
try:
resp = await self.client.get_metadata(self.id, "label")
except NonStreamableError as e:
logger.error(f"Error resolving Label: {e}")
return None
try:
meta = LabelMetadata.from_resp(resp, self.client.source)
except Exception as e:
logger.error(f"Error resolving Label: {e}")
return None
albums = [
PendingAlbum(album_id, self.client, self.config, self.db)
for album_id in meta.album_ids()
Expand Down
6 changes: 5 additions & 1 deletion streamrip/media/playlist.py
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,11 @@ async def resolve(self) -> Playlist | None:
)
return None

meta = PlaylistMetadata.from_resp(resp, self.client.source)
try:
meta = PlaylistMetadata.from_resp(resp, self.client.source)
except Exception as e:
logger.error(f"Error creating playlist: {e}")
return None
name = meta.name
parent = self.config.session.downloads.folder
folder = os.path.join(parent, clean_filepath(name))
Expand Down
Loading

0 comments on commit 527b52c

Please sign in to comment.