yandex-music-downloader/ymd/core.py

368 lines
12 KiB
Python

import datetime as dt
import re
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,
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 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"
class CoreTrackQuality(IntEnum):
LOW = 0
NORMAL = auto()
LOSSLESS = auto()
class LyricsFormat(LowercaseStrEnum):
NONE = auto()
TEXT = auto()
LRC = auto()
@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) -> Client:
client = Client(token)
client.request.set_timeout(timeout)
return client.init()
def full_title(obj: YandexMusicModel) -> str:
result = obj["title"]
if result is None:
return ""
if version := obj["version"]:
result += f" ({version})"
return result
def prepare_base_path(
path_pattern: Path, track: Track, unsafe_path: bool = False
) -> Path:
path_str = str(path_pattern)
album = None
artist = None
track_position = None
if albums := track.albums:
album = albums[0]
track_position = album.track_position
if artists := album.artists:
artist = artists[0]
if artist is None and (artists := track.artists):
artist = artists[0]
repl_dict: dict[str, Union[str, int, None]] = {
"#number-padded": str(track_position.index).zfill(len(str(album.track_count)))
if track_position and album
else None,
"#album-artist": artist.name if artist else None,
"#artist-id": artist.id if artist else None,
"#album-id": album.id if album else None,
"#track-id": track.id,
"#number": track_position.index if track_position else None,
"#artist": artist.name if artist else None,
"#title": full_title(track),
"#album": full_title(album) if album else None,
"#year": album.year if album else None,
}
for placeholder, replacement in repl_dict.items():
replacement = str(replacement)
if not unsafe_path:
clear_re = SAFE_PATH_CLEAR_RE
else:
clear_re = UNSAFE_PATH_CLEAR_RE
replacement = clear_re.sub("_", replacement)
path_str = path_str.replace(placeholder, replacement)
return Path(path_str)
def set_tags(
path: Path,
track: Track,
lyrics: Optional[str],
album_cover: Optional[AlbumCover],
compatibility_level: int,
) -> None:
album = track.albums[0] if track.albums else Album()
track_artists = [a.name for a in track.artists if a.name]
album_artists = [a.name for a in album.artists if a.name]
tag = mutagen.File(path, [MP3, MP4, FLAC]) # type: ignore
album_title = full_title(album)
track_title = full_title(track)
track_number = None
disc_number = None
if position := album.track_position:
track_number = position.index
disc_number = position.volume
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 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["rldt"] = iso8601_release_date
if 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 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 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 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,
):
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)
write_via_temporary_file(
track_data,
target_path,
temporary_file_hook=lambda tmp_path: set_tags(
tmp_path, track, text_lyrics, cover, compatibility_level
),
)
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:
temporary_file = target_path.parent / (
TEMPORARY_FILE_NAME_TEMPLATE.format(target_path.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