From 659d7e5cf9bd856a6b203a78ff08d6afe716da71 Mon Sep 17 00:00:00 2001
From: Lev Plyusnin <me@llistochek.ru>
Date: Sat, 3 May 2025 11:59:47 +0700
Subject: [PATCH] =?UTF-8?q?=D0=9A=D0=BE=D1=80=D1=80=D0=B5=D0=BA=D1=82?=
 =?UTF-8?q?=D0=BD=D0=BE=D0=B5=20=D1=80=D0=B0=D1=81=D1=88=D0=B8=D1=80=D0=B5?=
 =?UTF-8?q?=D0=BD=D0=B8=D0=B5=20=D1=84=D0=B0=D0=B9=D0=BB=D0=BE=D0=B2?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pyproject.toml |  2 +-
 ymd/api.py     | 49 ++++++++++++++++++++++++++++++-------------------
 ymd/cli.py     |  2 +-
 ymd/core.py    | 12 ++++++------
 4 files changed, 38 insertions(+), 27 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 73df93c..dcce542 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "yandex-music-downloader"
-version = "3.4.5b0"
+version = "3.4.5b1"
 description = "Загрузчик музыки с сервиса Яндекс.Музыка"
 requires-python = ">=3.9"
 readme = "README.md"
diff --git a/ymd/api.py b/ymd/api.py
index 8938fb4..2e153ec 100644
--- a/ymd/api.py
+++ b/ymd/api.py
@@ -1,13 +1,11 @@
 import base64
 import hashlib
 import hmac
-import itertools
 import random
 import time
 import typing
-from collections.abc import Iterable
 from dataclasses import dataclass
-from enum import Enum
+from enum import Enum, auto
 
 from Crypto.Cipher import AES
 from strenum import StrEnum
@@ -15,20 +13,33 @@ from yandex_music import Client, Track
 from yandex_music.utils.sign_request import DEFAULT_SIGN_KEY
 
 
+class Container(Enum):
+    FLAC = auto()
+    MP3 = auto()
+    MP4 = auto()
+
+
 class Codec(Enum):
-    FLAC = ["flac", "flac-mp4"]
-    MP3 = ["mp3"]
-    AAC = ["aac", "he-aac", "aac-mp4", "he-aac-mp4"]
+    FLAC = auto()
+    MP3 = auto()
+    AAC = auto()
 
-    @classmethod
-    def get_all_codecs(cls) -> Iterable[str]:
-        return itertools.chain.from_iterable(e.value for e in cls)
 
-    @classmethod
-    def from_codec_str(cls, codec_str: str) -> typing.Optional["Codec"]:
-        for codec in cls:
-            if any(e == codec_str for e in codec.value):
-                return codec
+@dataclass
+class FileFormat:
+    container: Container
+    codec: Codec
+
+
+FILE_FORMAT_MAPPING = {
+    "flac": FileFormat(Container.FLAC, Codec.FLAC),
+    "flac-mp4": FileFormat(Container.MP4, Codec.FLAC),
+    "mp3": FileFormat(Container.MP3, Codec.MP3),
+    "aac": FileFormat(Container.MP4, Codec.AAC),
+    "he-aac": FileFormat(Container.MP4, Codec.AAC),
+    "aac-mp4": FileFormat(Container.MP4, Codec.AAC),
+    "he-aac-mp4": FileFormat(Container.MP4, Codec.AAC),
+}
 
 
 class ApiTrackQuality(StrEnum):
@@ -40,7 +51,7 @@ class ApiTrackQuality(StrEnum):
 @dataclass
 class CustomDownloadInfo:
     quality: str
-    codec: Codec
+    file_format: FileFormat
     urls: list[str]
     decryption_key: str
     bitrate: int
@@ -54,7 +65,7 @@ def get_download_info(track: Track, quality: ApiTrackQuality) -> CustomDownloadI
         "ts": timestamp,
         "trackId": track.id,
         "quality": quality,
-        "codecs": ",".join(Codec.get_all_codecs()),
+        "codecs": ",".join(FILE_FORMAT_MAPPING.keys()),
         "transports": "encraw",
     }
     hmac_sign = hmac.new(
@@ -71,12 +82,12 @@ def get_download_info(track: Track, quality: ApiTrackQuality) -> CustomDownloadI
     resp = typing.cast(dict, resp)
     e = resp["download_info"]
     raw_codec = e["codec"]
-    codec = Codec.from_codec_str(raw_codec)
-    if codec is None:
+    file_format = FILE_FORMAT_MAPPING.get(raw_codec)
+    if file_format is None:
         raise ValueError(f"Unknown codec: {raw_codec}")
     return CustomDownloadInfo(
         quality=e["quality"],
-        codec=codec,
+        file_format=file_format,
         urls=e["urls"],
         bitrate=e["bitrate"],
         decryption_key=e.get("key"),
diff --git a/ymd/cli.py b/ymd/cli.py
index e125c68..fe233c9 100644
--- a/ymd/cli.py
+++ b/ymd/cli.py
@@ -302,7 +302,7 @@ def main():
 
         downloadable = core.to_downloadable_track(track, args.quality, save_path)
         bitrate = downloadable.download_info.bitrate
-        format_info = "[" + downloadable.download_info.codec.name
+        format_info = "[" + downloadable.download_info.file_format.codec.name
         if bitrate > 0:
             format_info += f" {bitrate}kbps"
         format_info += "]"
diff --git a/ymd/core.py b/ymd/core.py
index f29910a..8fcd46a 100644
--- a/ymd/core.py
+++ b/ymd/core.py
@@ -26,16 +26,16 @@ from mutagen.mp3 import MP3
 from mutagen.mp4 import MP4, MP4Cover
 from strenum import LowercaseStrEnum
 from yandex_music import (
+    Album,
     Client,
     Track,
     YandexMusicModel,
-    Album,
 )
 
 from ymd import api
 from ymd.api import (
     ApiTrackQuality,
-    Codec,
+    Container,
     CustomDownloadInfo,
     get_download_info,
 )
@@ -330,13 +330,13 @@ def to_downloadable_track(
         api_quality = ApiTrackQuality.LOSSLESS
 
     download_info = get_download_info(track, api_quality)
-    codec = download_info.codec
+    container = download_info.file_format.container
 
-    if codec == Codec.MP3:
+    if container == Container.MP3:
         suffix = ".mp3"
-    elif codec == Codec.AAC:
+    elif container == Container.MP4:
         suffix = ".m4a"
-    elif codec == Codec.FLAC:
+    elif container == Container.FLAC:
         suffix = ".flac"
     else:
         raise RuntimeError("Unknown codec")