From 1c2bf4f3117fe513cd95e82fea09326d2d75e296 Mon Sep 17 00:00:00 2001 From: Lev Plyusnin Date: Tue, 24 Dec 2024 17:39:56 +0700 Subject: [PATCH] =?UTF-8?q?=D0=9F=D0=BE=D0=B4=D0=B4=D0=B5=D1=80=D0=B6?= =?UTF-8?q?=D0=BA=D0=B0=20.png=20=D0=BE=D0=B1=D0=BB=D0=BE=D0=B6=D0=B5?= =?UTF-8?q?=D0=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 2 +- ymd/core.py | 49 ++++++++++++++++++++++++++++++++++++----------- ymd/mime_utils.py | 19 ++++++++++++++++++ 3 files changed, 58 insertions(+), 12 deletions(-) create mode 100644 ymd/mime_utils.py diff --git a/pyproject.toml b/pyproject.toml index 90f2164..5604c39 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" [project] name = "yandex-music-downloader" -version = "3.2.3b0" +version = "3.3.0b0" description = "Загрузчик музыки с сервиса Яндекс.Музыка" requires-python = ">=3.9" readme = "README.md" diff --git a/ymd/core.py b/ymd/core.py index b6cf0a2..69ca7d9 100644 --- a/ymd/core.py +++ b/ymd/core.py @@ -28,6 +28,7 @@ from strenum import LowercaseStrEnum from yandex_music import Client, DownloadInfo, Track, YandexMusicObject from ymd.api import get_lossless_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\-\'() ]+") @@ -56,6 +57,12 @@ class DownloadableTrack: track: Track +@dataclass +class AlbumCover: + data: bytes + mime_type: MimeType + + def init_client(token: str) -> Client: return Client(token).init() @@ -112,7 +119,7 @@ def set_tags( path: Path, track: Track, lyrics: Optional[str], - album_cover: Optional[bytes], + album_cover: Optional[AlbumCover], compatibility_level: int, ) -> None: album = track.albums[0] @@ -154,7 +161,12 @@ def set_tags( 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["APIC"] = APIC( + encoding=3, + mime=album_cover.mime_type.value, + type=3, + data=album_cover.data, + ) tag["WOAF"] = WOAF( encoding=3, @@ -183,7 +195,14 @@ def set_tags( if lyrics: tag["\xa9lyr"] = lyrics if album_cover: - tag["covr"] = [MP4Cover(album_cover, imageformat=MP4Cover.FORMAT_JPEG)] + 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 @@ -203,8 +222,8 @@ def set_tags( if album_cover is not None: pic = Picture() pic.type = PictureType.COVER_FRONT - pic.data = album_cover - pic.mime = "image/jpeg" + pic.data = album_cover.data + pic.mime = album_cover.mime_type.value tag.add_picture(pic) tag["comment"] = track_url else: @@ -218,12 +237,12 @@ def download_track( cover_resolution: int = DEFAULT_COVER_RESOLUTION, lyrics_format: LyricsFormat = LyricsFormat.NONE, embed_cover: bool = False, - covers_cache: Optional[dict[int, bytes]] = None, + 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, bytes], covers_cache) + covers_cache = typing.cast(dict[int, AlbumCover], covers_cache) target_path = track_info.path track = track_info.track client = track.client @@ -251,18 +270,26 @@ def download_track( 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_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 + cover = covers_cache[album_id] = album_cover else: - cover_path = target_path.parent / "cover.jpg" + 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(): - track.download_cover(str(cover_path), cover_size) + cover_path.write_bytes(album_cover.data) set_tags(target_path, track, text_lyrics, cover, compatibility_level) diff --git a/ymd/mime_utils.py b/ymd/mime_utils.py new file mode 100644 index 0000000..621dad1 --- /dev/null +++ b/ymd/mime_utils.py @@ -0,0 +1,19 @@ +from enum import Enum +from typing import Optional + + +class MimeType(Enum): + JPEG = "image/jpeg" + PNG = "image/png" + + +MAGIC_BYTES = ( + (MimeType.JPEG, bytes((0xFF, 0xD8, 0xFF))), + (MimeType.PNG, bytes((0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A))), +) + + +def guess_mime_type(data: bytes) -> Optional[MimeType]: + for mime_type, magic_bytes in MAGIC_BYTES: + if data.startswith(magic_bytes): + return mime_type