import datetime as dt import random import re import typing from dataclasses import dataclass 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 yandex_music import Client, DownloadInfo, Track, YandexMusicObject from ymd.api import get_lossless_info FILENAME_CLEAR_RE = re.compile(r"[^\w\-\'() ]+") DEFAULT_PATH_PATTERN = Path("#album-artist", "#album", "#number - #title") DEFAULT_COVER_RESOLUTION = 400 MIN_COMPATIBILITY_LEVEL = 0 MAX_COMPATIBILITY_LEVEL = 1 @dataclass class DownloadableTrack: url: str bitrate: int codec: str path: Path track: Track def init_client(token: str) -> Client: return Client(token).init() def full_title(obj: YandexMusicObject) -> Optional[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 = track.albums[0] artist = album.artists[0] track_position = album.track_position repl_dict: dict[str, Union[str, int, None]] = { "#number-padded": str(track_position.index).zfill(len(str(album.track_count))) if track_position else None, "#album-artist": album.artists[0].name, "#artist-id": artist.name, "#album-id": album.id, "#track-id": track.id, "#number": track_position.index if track_position else None, "#artist": artist.name, "#title": full_title(track), "#album": full_title(album), "#year": album.year, } for placeholder, replacement in repl_dict.items(): replacement = str(replacement) if not unsafe_path: replacement = FILENAME_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[bytes], compatibility_level: int, ) -> None: album = track.albums[0] 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 if album.release_date is not None: iso8601_release_date = dt.datetime.fromisoformat(album.release_date).astimezone( dt.timezone.utc ) iso8601_release_date = iso8601_release_date.strftime("%Y-%m-%d %H:%M:%S") release_year = None 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="image/jpeg", type=3, data=album_cover) 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 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 lyrics: tag["\xa9lyr"] = lyrics if album_cover: tag["covr"] = [MP4Cover(album_cover, imageformat=MP4Cover.FORMAT_JPEG)] 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 pic.mime = "image/jpeg" 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, add_lyrics: bool = False, embed_cover: bool = False, covers_cache: Optional[dict[int, bytes]] = 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, bytes], covers_cache) target_path = track_info.path track = track_info.track client = track.client assert client album = track.albums[0] client.request.download(track_info.url, str(target_path)) lyrics = None if add_lyrics and (lyrics_info := track.lyrics_info): if lyrics_info.has_available_text_lyrics: if track_lyrics := track.get_lyrics(format="TEXT"): lyrics = track_lyrics.fetch_lyrics() cover = None if track.cover_uri is not None: cover_size = f"{cover_resolution}x{cover_resolution}" if embed_cover: album_id = album.id if album_id and (cached_cover := covers_cache.get(album_id)): cover = cached_cover else: cover_bytes = track.download_cover_bytes(size=cover_size) if album_id: cover = covers_cache[album_id] = cover_bytes else: cover_path = target_path.parent / "cover.jpg" if not cover_path.is_file(): track.download_cover(str(cover_path), cover_size) set_tags(target_path, track, lyrics, cover, compatibility_level) def to_downloadable_track( track: Track, quality: int, base_path: Path ) -> DownloadableTrack: url: str codec: str bitrate: int codec: str if quality == 2: download_info = get_lossless_info(track) codec = download_info.codec url = random.choice(download_info.urls) bitrate = download_info.bitrate else: download_info = track.get_download_info(get_direct_links=True) download_info = [e for e in download_info if e.codec in ("mp3", "aac")] def sort_key(e: DownloadInfo) -> int | float: aac_multiplier = 1.5 bitrate = e.bitrate_in_kbps if bitrate <= 192: aac_multiplier = 0.5 if e.codec == "aac": bitrate *= aac_multiplier return bitrate download_info.sort( key=sort_key, reverse=quality == 0, ) target_info = download_info[-1] url = typing.cast(str, target_info.direct_link) bitrate = target_info.bitrate_in_kbps codec = target_info.codec if codec == "mp3": suffix = ".mp3" elif codec == "aac" or codec == "he-aac": suffix = ".m4a" elif codec == "flac": suffix = ".flac" else: raise RuntimeError("Unknown codec") target_path = base_path.with_suffix(suffix) return DownloadableTrack( url=url, track=track, bitrate=bitrate, codec=codec, path=target_path, )