Create-knowledge-find-decode/creation_knowledge/media.py

"""媒体本地落盘：把图文帖的图片、视频帖的整段 mp4 存到 data/<run_id>/<platform>/<post_id>/。

解耦设计：仅当 settings.data_dir 非空时调用；视频落盘在 extract_video 内（复用其内存
字节），这里只管图片下载 + 卡片 url 改写为本地 /data 路径。
"""
from __future__ import annotations

import datetime
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

from creation_knowledge.integrations.video_extract import _default_download
from core.models import Post

logger = logging.getLogger(__name__)

Downloader = Callable[[str, str], bytes]


def run_id() -> str:
    """运行批次 id，默认时间戳到秒。"""
    return datetime.datetime.now().strftime("%Y%m%d-%H%M%S")


@dataclass
class MediaTarget:
    dir: Path           # 本地目录 data/<run_id>/<platform>/<post_id>/
    public_base: str    # 浏览器前缀 /data/<run_id>/<platform>/<post_id>


def media_target(data_dir: str, batch: str, platform: str, post_id: str) -> MediaTarget:
    rel = f"{batch}/{platform}/{post_id}"
    return MediaTarget(dir=Path(data_dir) / rel, public_base=f"/data/{rel}")


def _ext_from_bytes(b: bytes) -> str:
    """按文件头嗅探扩展名（不依赖 Content-Type，便于注入式测试）。"""
    if b[:4] == b"\x89PNG":
        return ".png"
    if b[:3] == b"GIF":
        return ".gif"
    if b[:4] == b"RIFF" and b[8:12] == b"WEBP":
        return ".webp"
    return ".jpg"


def save_images(
    post: Post,
    media_dir: Path,
    public_base: str,
    *,
    downloader: Downloader = _default_download,
) -> None:
    """下载图文帖每张远程图片 → media_dir/image_<index>.<ext>，并把 card.url 改写为本地路径。

    单张失败只告警跳过（不影响其余卡片与整条流水线）。
    """
    media_dir = Path(media_dir)
    for card in post.cards:
        if card.kind != "image" or not card.url or not card.url.startswith("http"):
            continue
        try:
            data = downloader(card.url, post.platform)
        except Exception as exc:  # 单张图失败不致命
            logger.warning("图片下载失败 idx=%s %s: %s", card.index, card.url, exc)
            continue
        if not data:
            continue
        media_dir.mkdir(parents=True, exist_ok=True)
        fname = f"image_{card.index}{_ext_from_bytes(data)}"
        (media_dir / fname).write_bytes(data)
        card.url = f"{public_base}/{fname}"