NeteaseUrl/music_downloader/tagging.py

# -*- coding: utf-8 -*-
"""音频标签与封面写入（mutagen 纯 Python 实现，不需要 ffmpeg 命令）。

为什么不再用 ffmpeg CLI：本工程的元数据写入只做「改标签」，不做任何重编码，
而 ffmpeg 命令行必须出现**重复的** `-metadata` / `-map` 参数；ffmpeg-python 用
`**kwargs` 字典建模这些选项，结构上无法表达重复键（实测把列表值字符串化成
`"['0', '1:v']"`），绕一圈仍需手写命令拼接。改用 mutagen 后整段变成纯 Python
调用，部署镜像也不必再装 ffmpeg。

覆盖容器：
- MP3        ID3v2.3（TIT2 / TPE1 / TPE2 / TALB / TDRC / TRCK / USLT / APIC）
- FLAC       Vorbis comment（title / artist / albumartist / …）+ Picture 块
- M4A / MP4  iTunes ilst（©nam / ©ART / aART / ©alb / ©day / trkn / ©lyr）+ covr
- OGG / Opus Vorbis comment + METADATA_BLOCK_PICTURE（base64 编码的 Picture 块）

写入一律走 mutagen 的原子替换（临时文件 + os.replace），中途失败不会留下半截文件；
且是「就地改帧」，原有未触及的标签（TXXX 等）会被保留。
"""

from __future__ import annotations

import base64
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from ._log import logger
from .models import METADATA_FORMATS, MusicInfo

# 封面类型：ID3 APIC / FLAC Picture 的 picture type 3 = Cover (front)
_PICTURE_TYPE_COVER_FRONT = 3
_COVER_DESC = 'Cover (front)'
_PNG_MAGIC = b'\x89PNG\r\n\x1a\n'

# LRC 时间标签（[mm:ss] / [mm:ss.xx] / [mm:ss:xx]）
_LRC_TIME_TAG = re.compile(r'\[\d{1,2}:\d{1,2}([.:]\d+)?\]')


@dataclass
class TagSet:
    """与容器无关的标签中间表示；None 表示本次不写该字段。"""

    title: str | None = None
    artist: str | None = None
    album: str | None = None
    album_artist: str | None = None
    date: str | None = None
    track: int | None = None
    lyrics: str | None = None


def build_tags(info: MusicInfo, fields: list[str] | None) -> TagSet:
    """按前端勾选的字段构造标签集合；fields 为 None 表示全部写入。"""
    def want(*names: str) -> bool:
        return fields is None or any(n in fields for n in names)

    tags = TagSet()
    if want('title', 'name'):
        tags.title = info.name or None
    if want('artist'):
        tags.artist = info.artists or None
        tags.album_artist = info.artists or None
    if want('album'):
        tags.album = info.album or None
    if want('date') and info.publish_year:
        tags.date = str(info.publish_year)
    if want('lyrics') and info.lyric:
        # 写入纯歌词：去掉 LRC 时间标签
        tags.lyrics = _LRC_TIME_TAG.sub('', info.lyric).strip() or None
    if want('track') and info.track_number > 0:
        # 只写当前曲号；没有真实专辑总数时不虚构 tracktotal
        tags.track = info.track_number
    return tags


def write_tags(path: Path, tags: TagSet, cover: bytes | None = None) -> bool:
    """把标签与封面写入音频文件；失败仅记日志并返回 False（不影响交付主流程）。"""
    ext = path.suffix.lower()
    writer = _WRITERS.get(ext)
    if writer is None:
        logger.warning(f"跳过元数据写入: 不支持的格式 {ext}")
        return False
    if not path.exists():
        logger.warning(f"跳过元数据写入: 文件不存在 {path}")
        return False
    try:
        writer(path, tags, cover)
        return True
    except Exception as e:
        logger.warning(f"写入元数据失败（{ext}）: {e}")
        return False


# ---------- 各容器实现 ----------

def _write_mp3(path: Path, tags: TagSet, cover: bytes | None) -> None:
    from mutagen.id3 import (
        APIC,
        ID3,
        TALB,
        TDRC,
        TIT2,
        TPE1,
        TPE2,
        TRCK,
        USLT,
        ID3NoHeaderError,
    )

    try:
        audio = ID3(str(path))
    except ID3NoHeaderError:
        audio = ID3()                      # 无 ID3 头的裸音频：新建标签区
    if tags.title:
        audio['TIT2'] = TIT2(encoding=3, text=[tags.title])
    if tags.artist:
        audio['TPE1'] = TPE1(encoding=3, text=[tags.artist])
    if tags.album_artist:
        audio['TPE2'] = TPE2(encoding=3, text=[tags.album_artist])
    if tags.album:
        audio['TALB'] = TALB(encoding=3, text=[tags.album])
    if tags.date:
        audio['TDRC'] = TDRC(encoding=3, text=[tags.date])
    if tags.track:
        audio['TRCK'] = TRCK(encoding=3, text=[str(tags.track)])
    if tags.lyrics:
        audio['USLT'] = USLT(encoding=3, lang='eng', desc='', text=tags.lyrics)
    if cover:
        audio['APIC'] = APIC(encoding=3, mime=_cover_mime(cover),
                             type=_PICTURE_TYPE_COVER_FRONT,
                             desc=_COVER_DESC, data=cover)
    # ID3v2.3 与 ffmpeg 时代的 -id3v2_version 3 保持一致（mutagen 会自动把
    # TDRC 降级为 TYER/TDAT）；不写 v1 标签以免在文件尾追加 128 字节
    audio.save(str(path), v2_version=3)


def _write_flac(path: Path, tags: TagSet, cover: bytes | None) -> None:
    from mutagen.flac import FLAC

    audio = FLAC(str(path))
    _set_vorbis_comments(audio, tags)
    if cover:
        audio.clear_pictures()
        audio.add_picture(_make_picture(cover))
    audio.save()


def _write_m4a(path: Path, tags: TagSet, cover: bytes | None) -> None:
    from mutagen.mp4 import MP4, MP4Cover

    audio = MP4(str(path))
    mapping = {
        '\xa9nam': tags.title,
        '\xa9ART': tags.artist,
        'aART': tags.album_artist,
        '\xa9alb': tags.album,
        '\xa9day': tags.date,
        '\xa9lyr': tags.lyrics,
    }
    for key, value in mapping.items():
        if value:
            audio[key] = [value]
    if tags.track:
        audio['trkn'] = [(tags.track, 0)]   # 不虚构 tracktotal
    if cover:
        image_format = (MP4Cover.FORMAT_PNG if cover.startswith(_PNG_MAGIC)
                        else MP4Cover.FORMAT_JPEG)
        audio['covr'] = [MP4Cover(cover, imageformat=image_format)]
    audio.save()


def _write_ogg(path: Path, tags: TagSet, cover: bytes | None) -> None:
    _write_ogg_stream(path, tags, cover, opus=False)


def _write_opus(path: Path, tags: TagSet, cover: bytes | None) -> None:
    _write_ogg_stream(path, tags, cover, opus=True)


def _write_ogg_stream(path: Path, tags: TagSet, cover: bytes | None, *,
                      opus: bool) -> None:
    from mutagen.oggopus import OggOpus
    from mutagen.oggvorbis import OggVorbis

    audio: Any = OggOpus(str(path)) if opus else OggVorbis(str(path))
    _set_vorbis_comments(audio, tags)
    if cover:
        # Ogg 系把封面按 METADATA_BLOCK_PICTURE 规范存成 base64 的 Picture 块
        audio['metadata_block_picture'] = [
            base64.b64encode(_make_picture(cover).write()).decode('ascii')]
    audio.save()


# ---------- 小工具 ----------

def _cover_mime(cover: bytes) -> str:
    return 'image/png' if cover.startswith(_PNG_MAGIC) else 'image/jpeg'


def _make_picture(cover: bytes) -> Any:
    from mutagen.flac import Picture

    pic = Picture()
    pic.type = _PICTURE_TYPE_COVER_FRONT
    pic.mime = _cover_mime(cover)
    pic.desc = _COVER_DESC
    pic.data = cover
    return pic


def _set_vorbis_comments(audio: Any, tags: TagSet) -> None:
    """FLAC 与 Ogg 共用 Vorbis comment 字段名。"""
    mapping = {
        'title': tags.title,
        'artist': tags.artist,
        'albumartist': tags.album_artist,
        'album': tags.album,
        'date': tags.date,
    }
    for key, value in mapping.items():
        if value:
            audio[key] = [value]
    if tags.track:
        audio['tracknumber'] = [str(tags.track)]
    if tags.lyrics:
        audio['lyrics'] = [tags.lyrics]


_WRITERS = {
    '.mp3': _write_mp3,
    '.flac': _write_flac,
    '.m4a': _write_m4a,
    '.mp4': _write_m4a,
    '.ogg': _write_ogg,
    '.opus': _write_opus,
}

# 与 models 的类型表保持同源：新增容器必须两边同时登记
assert set(_WRITERS) == set(METADATA_FORMATS), '标签写入表与 METADATA_FORMATS 不一致'