yandex-music-downloader/ymd/core.py

705 lines
22 KiB
Python

import datetime as dt
import hashlib
import re
import time
import typing
from collections.abc import Callable
from dataclasses import dataclass
from enum import IntEnum, auto
from pathlib import Path
from typing import Optional, Union
import mutagen
from mutagen.flac import FLAC, Picture
from mutagen.id3._frames import (
APIC,
TALB,
TCON,
TDRC,
TIT2,
TPE1,
TPE2,
TPOS,
TRCK,
USLT,
WOAF,
)
from mutagen.id3._specs import ID3TimeStamp, PictureType
from mutagen.mp3 import MP3
from mutagen.mp4 import MP4, MP4Cover
from strenum import LowercaseStrEnum
from yandex_music import (
Album,
Client,
Track,
YandexMusicModel,
)
from yandex_music.exceptions import NetworkError
from ymd import api
from ymd.api import (
ApiTrackQuality,
Container,
CustomDownloadInfo,
get_download_info,
)
from ymd.mime_utils import MimeType, guess_mime_type
UNSAFE_PATH_CLEAR_RE = re.compile(r"[/\\]+")
SAFE_PATH_CLEAR_RE = re.compile(r"([^\w\-\'() ]|^\s+|\s+$)")
DEFAULT_PATH_PATTERN = Path("#album-artist", "#album", "#number-#title")
DEFAULT_COVER_RESOLUTION = 400
MIN_COMPATIBILITY_LEVEL = 0
MAX_COMPATIBILITY_LEVEL = 1
AUDIO_FILE_SUFFIXES = {".mp3", ".flac", ".m4a"}
TEMPORARY_FILE_NAME_TEMPLATE = ".yandex-music-downloader.{}.tmp"
MAX_FILE_NAME_LENGTH_WITHOUT_SUFFIX = 255 - max(
len(suffix) for suffix in AUDIO_FILE_SUFFIXES
)
class CoreTrackQuality(IntEnum):
LOW = 0
NORMAL = auto()
LOSSLESS = auto()
class LyricsFormat(LowercaseStrEnum):
NONE = auto()
TEXT = auto()
LRC = auto()
CONTAINER_MUTAGEN_MAPPING: dict[Container, type[mutagen.FileType]] = { # type: ignore
Container.MP3: MP3,
Container.FLAC: FLAC,
Container.MP4: MP4,
}
@dataclass
class DownloadableTrack:
download_info: CustomDownloadInfo
path: Path
track: Track
@dataclass
class AlbumCover:
data: bytes
mime_type: MimeType
def init_client(
token: str, timeout: int, max_try_count: int, retry_delay: int
) -> Client:
assert timeout > 0
assert max_try_count >= 0
assert retry_delay >= 0
client = Client(token)
client.request.set_timeout(timeout)
original_wrapper = client.request._request_wrapper
def retry_wrapper(*args, **kwargs):
try_count = 0
while True:
try:
return original_wrapper(*args, **kwargs)
except NetworkError as error:
if max_try_count == 0 or try_count < max_try_count:
try_count += 1
time.sleep(retry_delay)
continue
raise error
client.request._request_wrapper = retry_wrapper
return client.init()
def full_title(obj: YandexMusicModel) -> str:
"""Return title with optional version; be tolerant to missing keys."""
title = getattr(obj, "title", None)
if title is None:
try:
title = obj["title"]
except Exception: # noqa: BLE001
title = None
if not title:
return ""
version = getattr(obj, "version", None)
if version is None:
try:
version = obj["version"]
except Exception: # noqa: BLE001
version = None
if version:
title = f"{title} ({version})"
return title
def sanitize_path_component(text: str, unsafe_path: bool) -> str:
"""Return path-safe component."""
if unsafe_path:
return UNSAFE_PATH_CLEAR_RE.sub("_", text)
text = text.strip()
text = SAFE_PATH_CLEAR_RE.sub("_", text)
text = re.sub(r"\s+", "_", text)
text = re.sub(r"_+", "_", text)
return text
def choose_first(*values: Optional[Union[str, int]]) -> str:
for candidate in values:
if candidate is None:
continue
text = str(candidate).strip()
if text:
return text
return ""
def prepare_base_path(
path_pattern: Path,
track: Track,
unsafe_path: bool = False,
fallback_number: Optional[int] = None,
fallback_total: Optional[int] = None,
override_title: Optional[str] = None,
override_album: Optional[str] = None,
override_album_artist: Optional[str] = None,
override_track_artist: Optional[str] = None,
override_track_number: Optional[int] = None,
) -> Path:
path_str = str(path_pattern)
album = None
album_artist = None
track_artist = None
track_position = None
if albums := track.albums:
album = albums[0]
if artists := album.artists:
album_artist = artists[0]
# позиция трека корректнее брать из самого трека
track_position = getattr(track, "track_position", None) or track_position
if artists := track.artists:
track_artist = artists[0]
track_id = getattr(track, "id", None)
track_title = choose_first(
override_title,
full_title(track),
f"track_{track_id}" if track_id is not None else None,
f"track_{fallback_number}" if fallback_number is not None else None,
"track",
)
album_title = choose_first(
override_album,
full_title(album) if album else None,
str(album.id) if album and album.id is not None else None,
"album",
)
album_artist_name = choose_first(
override_album_artist,
album_artist.name if album_artist else None,
"artist",
)
track_artist_name = choose_first(
override_track_artist,
track_artist.name if track_artist else None,
album_artist_name,
"artist",
)
pad_width = None
if track_position and album and album.track_count:
pad_width = max(2, len(str(album.track_count)))
elif fallback_total:
pad_width = max(2, len(str(fallback_total)))
elif fallback_number:
pad_width = max(2, len(str(fallback_number)))
number = (
override_track_number
if override_track_number is not None
else (track_position.index if track_position else fallback_number)
)
number_padded = (
str(number).zfill(pad_width) if number is not None and pad_width is not None else number
)
number_value = number_padded if number_padded is not None else number
repl_dict: dict[str, Union[str, int, None]] = {
"#number-padded": number_padded,
"#album-artist": album_artist_name,
"#track-artist": track_artist_name,
"#artist-id": track_artist.id if track_artist else None,
"#album-id": album.id if album else None,
"#track-id": track.id,
"#number": number_value,
"#title": track_title,
"#album": album_title,
"#year": album.year if album else None,
}
for placeholder, replacement in repl_dict.items():
if replacement is None:
replacement = ""
replacement = sanitize_path_component(str(replacement), unsafe_path)
if not replacement:
replacement = sanitize_path_component("untitled", unsafe_path)
path_str = path_str.replace(placeholder, replacement)
path = Path(path_str)
if not unsafe_path:
path = Path(
*(
re.sub(r"_+", "_", part.replace(" ", "_"))
for part in path.parts
)
)
trimmed_parts = [
part
if len(part) <= MAX_FILE_NAME_LENGTH_WITHOUT_SUFFIX
else part[:MAX_FILE_NAME_LENGTH_WITHOUT_SUFFIX]
for part in path.parts
]
return Path(*trimmed_parts)
def set_tags(
path: Path,
track: Track,
container: Container,
lyrics: Optional[str],
album_cover: Optional[AlbumCover],
compatibility_level: int,
source_tags: Optional[dict[str, Optional[Union[str, int]]]] = None,
) -> None:
file_type = CONTAINER_MUTAGEN_MAPPING.get(container)
if file_type is None:
raise ValueError(f"Unknown container: {container}")
source_tags = source_tags or {}
def as_list(value: Optional[Union[str, list[str], tuple[str, ...]]]) -> list[str]:
if value is None:
return []
if isinstance(value, (list, tuple, set)):
return [str(v).strip() for v in value if str(v).strip()]
text = str(value).strip()
return [text] if text else []
tag = file_type(path)
album = track.albums[0] if track.albums else Album()
album_title = choose_first(
source_tags.get("album"),
full_title(album),
str(album.id) if getattr(album, "id", None) is not None else None,
)
track_title = choose_first(
source_tags.get("title"),
full_title(track),
f"track_{track.id}" if getattr(track, "id", None) is not None else None,
"track",
)
track_artists = [a.name for a in track.artists if a.name]
if not track_artists:
track_artists = as_list(source_tags.get("track_artist"))
album_artists = [a.name for a in album.artists if a.name]
if not album_artists:
album_artists = as_list(source_tags.get("album_artist"))
if not track_artists:
track_artists = album_artists
genre = None
if album.genre:
genre = album.genre
track_number = None
disc_number = None
position = getattr(track, "track_position", None) or album.track_position
if position:
track_number = position.index
disc_number = position.volume
if track_number is None and (tag_track := source_tags.get("track_number")):
try:
track_number = int(tag_track)
except (TypeError, ValueError):
track_number = None
iso8601_release_date = None
release_year: Optional[str] = None
if album.release_date is not None:
iso8601_release_date = dt.datetime.fromisoformat(album.release_date).astimezone(
dt.timezone.utc
)
release_year = str(iso8601_release_date.year)
iso8601_release_date = iso8601_release_date.strftime("%Y-%m-%d %H:%M:%S")
if year := album.year:
release_year = str(year)
track_url = f"https://music.yandex.ru/album/{album.id}/track/{track.id}"
if isinstance(tag, MP3):
tag["TIT2"] = TIT2(encoding=3, text=track_title)
tag["TALB"] = TALB(encoding=3, text=album_title)
tag["TPE1"] = TPE1(encoding=3, text=track_artists)
tag["TPE2"] = TPE2(encoding=3, text=album_artists)
if tdrc_text := iso8601_release_date or release_year:
tag["TDRC"] = TDRC(encoding=3, text=[ID3TimeStamp(tdrc_text)])
if track_number:
tag["TRCK"] = TRCK(encoding=3, text=str(track_number))
if disc_number:
tag["TPOS"] = TPOS(encoding=3, text=str(disc_number))
if genre:
tag["TCON"] = TCON(encoding=3, text=genre)
if lyrics:
tag["USLT"] = USLT(encoding=3, text=lyrics)
if album_cover:
tag["APIC"] = APIC(
encoding=3,
mime=album_cover.mime_type.value,
type=3,
data=album_cover.data,
)
tag["WOAF"] = WOAF(
encoding=3,
text=track_url,
)
elif isinstance(tag, MP4):
tag["\xa9nam"] = track_title
tag["\xa9alb"] = album_title
artists_value = track_artists
album_artists_value = album_artists
if compatibility_level == 1:
artists_value = "; ".join(track_artists)
album_artists_value = "; ".join(album_artists)
tag["\xa9ART"] = artists_value
tag["aART"] = album_artists_value
if iso8601_release_date is not None:
tag["\xa9day"] = iso8601_release_date
elif release_year is not None:
tag["\xa9day"] = release_year
if track_number:
tag["trkn"] = [(track_number, 0)]
if disc_number:
tag["disk"] = [(disc_number, 0)]
if genre:
tag["\xa9gen"] = genre
if lyrics:
tag["\xa9lyr"] = lyrics
if album_cover:
mime_mp4_dict = {
MimeType.JPEG: MP4Cover.FORMAT_JPEG,
MimeType.PNG: MP4Cover.FORMAT_PNG,
}
mp4_image_format = mime_mp4_dict.get(album_cover.mime_type)
if mp4_image_format is None:
raise RuntimeError("Unsupported cover type")
tag["covr"] = [MP4Cover(album_cover.data, imageformat=mp4_image_format)]
tag["\xa9cmt"] = track_url
elif isinstance(tag, FLAC):
tag["title"] = track_title
tag["album"] = album_title
tag["artist"] = track_artists
tag["albumartist"] = album_artists
if date_text := iso8601_release_date or release_year:
tag["date"] = date_text
if track_number:
tag["tracknumber"] = str(track_number)
if disc_number:
tag["discnumber"] = str(disc_number)
if genre:
tag["genre"] = genre
if lyrics:
tag["lyrics"] = lyrics
if album_cover is not None:
pic = Picture()
pic.type = PictureType.COVER_FRONT
pic.data = album_cover.data
pic.mime = album_cover.mime_type.value
tag.add_picture(pic)
tag["comment"] = track_url
else:
raise RuntimeError("Unknown file format")
tag.save()
def extract_basic_tags(path: Path, container: Container) -> dict[str, Optional[Union[str, int]]]:
"""Read minimal tags to be used for renaming."""
title = album = album_artist = track_artist = None
track_number: Optional[int] = None
tag = CONTAINER_MUTAGEN_MAPPING[container](path)
def first_text(value) -> Optional[str]:
if value is None:
return None
if isinstance(value, list):
if not value:
return None
value = value[0]
if hasattr(value, "text"):
value = value.text
if isinstance(value, list):
value = value[0] if value else None
if isinstance(value, bytes):
try:
value = value.decode("utf-8", "ignore")
except Exception:
return None
return str(value) if value else None
if isinstance(tag, MP3):
title = first_text(tag.get("TIT2"))
album = first_text(tag.get("TALB"))
track_artist = first_text(tag.get("TPE1"))
album_artist = first_text(tag.get("TPE2"))
if trck := first_text(tag.get("TRCK")):
try:
track_number = int(str(trck).split("/")[0])
except ValueError:
track_number = None
elif isinstance(tag, MP4):
title = first_text(tag.tags.get("\xa9nam") if tag.tags else None)
album = first_text(tag.tags.get("\xa9alb") if tag.tags else None)
track_artist = first_text(tag.tags.get("\xa9ART") if tag.tags else None)
album_artist = first_text(tag.tags.get("aART") if tag.tags else None)
if trkn := tag.tags.get("trkn") if tag.tags else None:
try:
track_number = int(trkn[0][0])
except Exception:
track_number = None
elif isinstance(tag, FLAC):
def get_flac(key: str) -> Optional[str]:
values = tag.get(key)
if not values:
return None
return values[0]
title = get_flac("title")
album = get_flac("album")
artists_val = get_flac("artist")
if artists_val:
track_artist = artists_val
album_artist = get_flac("albumartist")
if trck := get_flac("tracknumber"):
try:
track_number = int(str(trck).split("/")[0])
except ValueError:
track_number = None
return {
"title": title,
"album": album,
"album_artist": album_artist,
"track_artist": track_artist,
"track_number": track_number,
}
def retarget_with_tags(
path_pattern: Path,
track: Track,
container: Container,
current_path: Path,
unsafe_path: bool,
fallback_number: Optional[int],
fallback_total: Optional[int],
override_tags: Optional[dict[str, Optional[Union[str, int]]]] = None,
) -> Path:
tags: dict[str, Optional[Union[str, int]]] = {}
if override_tags:
tags.update(override_tags)
try:
extracted = extract_basic_tags(current_path, container)
except Exception:
extracted = {}
for key, value in extracted.items():
if key not in tags or not tags[key]:
tags[key] = value
new_base = prepare_base_path(
path_pattern,
track,
unsafe_path=unsafe_path,
fallback_number=fallback_number,
fallback_total=fallback_total,
override_title=tags.get("title"),
override_album=tags.get("album"),
override_album_artist=tags.get("album_artist"),
override_track_artist=tags.get("track_artist"),
override_track_number=tags.get("track_number"),
)
new_target = new_base.with_suffix(current_path.suffix)
try:
same = new_target.resolve() == current_path.resolve()
except FileNotFoundError:
same = False
if same:
return current_path
new_target.parent.mkdir(parents=True, exist_ok=True)
current_path.replace(new_target)
return new_target
def download_track(
track_info: DownloadableTrack,
cover_resolution: int = DEFAULT_COVER_RESOLUTION,
lyrics_format: LyricsFormat = LyricsFormat.NONE,
embed_cover: bool = False,
covers_cache: Optional[dict[int, AlbumCover]] = None,
compatibility_level: int = 1,
path_pattern: Optional[Path] = None,
unsafe_path: bool = False,
fallback_number: Optional[int] = None,
fallback_total: Optional[int] = None,
rename_using_tags: bool = False,
):
if embed_cover and covers_cache is None:
raise RuntimeError("covers_cache isn't provided")
covers_cache = typing.cast(dict[int, AlbumCover], covers_cache)
target_path = track_info.path
track = track_info.track
client = typing.cast(Client, track.client)
assert client
text_lyrics = None
if lyrics_format != LyricsFormat.NONE and (lyrics_info := track.lyrics_info):
if lyrics_format == LyricsFormat.LRC and lyrics_info.has_available_sync_lyrics:
lrc_path = target_path.with_suffix(".lrc")
if not lrc_path.is_file() and (
track_lyrics := track.get_lyrics(format_="LRC")
):
lyrics = track_lyrics.fetch_lyrics()
write_via_temporary_file(lyrics.encode("utf-8"), lrc_path)
elif lyrics_info.has_available_text_lyrics:
if track_lyrics := track.get_lyrics(format_="TEXT"):
text_lyrics = track_lyrics.fetch_lyrics()
cover = None
if track.cover_uri is not None:
if cover_resolution == -1:
cover_size = "orig"
else:
cover_size = f"{cover_resolution}x{cover_resolution}"
cover_bytes = track.download_cover_bytes(size=cover_size)
mime_type = guess_mime_type(cover_bytes)
if mime_type is None:
raise RuntimeError("Unknown cover mime type")
album_cover = AlbumCover(data=cover_bytes, mime_type=mime_type)
if embed_cover:
album = track.albums[0] if track.albums else Album()
if album.id and (cached_cover := covers_cache.get(album.id)):
cover = cached_cover
else:
if album.id:
cover = covers_cache[album.id] = album_cover
else:
mime_suffix_dict = {MimeType.JPEG: ".jpg", MimeType.PNG: ".png"}
file_suffix = mime_suffix_dict.get(album_cover.mime_type)
if file_suffix is None:
raise RuntimeError("Unknown mime type")
cover_path = target_path.parent / ("cover" + file_suffix)
if not cover_path.is_file():
write_via_temporary_file(album_cover.data, cover_path)
download_info = track_info.download_info
track_data = api.download_track(client, download_info)
source_tags: dict[str, Optional[Union[str, int]]] = {}
def tag_hook(tmp_path: Path) -> None:
nonlocal source_tags
try:
source_tags = extract_basic_tags(tmp_path, download_info.file_format.container)
except Exception:
source_tags = {}
set_tags(
tmp_path,
track,
download_info.file_format.container,
text_lyrics,
cover,
compatibility_level,
source_tags=source_tags,
)
write_via_temporary_file(
track_data,
target_path,
temporary_file_hook=tag_hook,
)
if rename_using_tags and path_pattern is not None:
try:
renamed_path = retarget_with_tags(
path_pattern=path_pattern,
track=track,
container=download_info.file_format.container,
current_path=target_path,
unsafe_path=unsafe_path,
fallback_number=fallback_number,
fallback_total=fallback_total,
override_tags=source_tags if source_tags else None,
)
target_path = renamed_path
except Exception:
# Renaming is non-critical; ignore failures to keep download intact
pass
def to_downloadable_track(
track: Track, quality: CoreTrackQuality, base_path: Path
) -> DownloadableTrack:
api_quality = ApiTrackQuality.NORMAL
if quality == CoreTrackQuality.LOW:
api_quality = ApiTrackQuality.LOW
elif quality == CoreTrackQuality.NORMAL:
api_quality = ApiTrackQuality.NORMAL
elif quality == CoreTrackQuality.LOSSLESS:
api_quality = ApiTrackQuality.LOSSLESS
download_info = get_download_info(track, api_quality)
container = download_info.file_format.container
if container == Container.MP3:
suffix = ".mp3"
elif container == Container.MP4:
suffix = ".m4a"
elif container == Container.FLAC:
suffix = ".flac"
else:
raise RuntimeError("Unknown codec")
target_path = str(base_path) + suffix
return DownloadableTrack(
download_info=download_info,
track=track,
path=Path(target_path),
)
def write_via_temporary_file(
data: bytes,
target_path: Path,
temporary_file_hook: Optional[Callable[[Path], None]] = None,
) -> Path:
target_name = hashlib.sha256(target_path.name.encode()).hexdigest()
temporary_file = target_path.parent / (
TEMPORARY_FILE_NAME_TEMPLATE.format(target_name)
)
try:
temporary_file.write_bytes(data)
if temporary_file_hook is not None:
temporary_file_hook(temporary_file)
except InterruptedError as e:
temporary_file.unlink()
raise e
temporary_file.rename(target_path)
return target_path