Source code for remote_compression.probe

"""Local media inspection through ffprobe."""

import json
import subprocess
from dataclasses import dataclass, field
from pathlib import Path




@dataclass(frozen=True, slots=True)
class SubtitleStream:
    """
    A subtitle stream as reported by ffprobe.

    Attributes
    ----------
    index: :class:`int`
        Index of the stream *among subtitle streams* (0-based, order of appearance),
        i.e. the ``i`` of ffmpeg specifiers like ``0:s:i``.
    codec_name: :class:`str`
        ffprobe codec name (e.g. ``subrip``, ``mov_text``, ``hdmv_pgs_subtitle``).
    """

    index: int
    codec_name: str





@dataclass(frozen=True, slots=True)
class MediaInfo:
    """
    Result of a successful probe: the first genuine video stream plus subtitles.

    Attributes
    ----------
    width: :class:`int`
        Video width in pixels.
    height: :class:`int`
        Video height in pixels.
    codec_name: :class:`str`
        ffprobe codec name of the video stream (e.g. ``h264``, ``hevc``).
    subtitles: :class:`tuple` of :class:`SubtitleStream`
        Subtitle streams, in order of appearance.
    duration: :class:`float`, optional
        Container duration in seconds, if known. Used as progress-bar total.
    field_order: :class:`str`, optional
        ffprobe field order (``progressive``, ``tt``, ``bb``...) when reported;
        drives automatic deinterlacing.
    """

    width: int
    height: int
    codec_name: str
    subtitles: tuple[SubtitleStream, ...] = field(default=())
    duration: float | None = None
    field_order: str | None = None



def _rotated_quarter_turn(stream):
    """True when display-matrix side data rotates the frame by +/-90 degrees."""
    for side_data in stream.get("side_data_list", []):
        rotation = side_data.get("rotation")
        if rotation is None:
            continue
        try:
            return abs(int(rotation)) % 180 == 90
        except (TypeError, ValueError):
            return False
    return False




def probe(file):
    """
    Inspect a local media file with ffprobe.

    Parameters
    ----------
    file: :class:`~pathlib.Path` or :class:`str`
        File location.

    Returns
    -------
    :class:`MediaInfo` or None
        Extracted information, or None whenever the file cannot be probed:
        ffprobe missing, non-zero exit, empty or invalid JSON output, or no
        genuine video stream (cover art, aka ``attached_pic``, does not count).

    Notes
    -----
    The command is passed as an argument list (never through a shell), so this
    works identically on Windows and POSIX.

    Width and height are *display* dimensions: when the stream carries a
    display-matrix rotation of +/-90 degrees (phone videos), the coded
    dimensions are swapped accordingly.
    """
    argv = [
        "ffprobe",
        "-v", "error",
        "-print_format", "json",
        "-show_streams",
        "-show_format",
        str(Path(file)),
    ]
    try:
        result = subprocess.run(argv, capture_output=True, check=False)
    except OSError:
        return None
    if result.returncode != 0 or not result.stdout.strip():
        return None
    try:
        data = json.loads(result.stdout)
    except json.JSONDecodeError:
        return None
    streams = data.get("streams", [])
    video = next(
        (
            s
            for s in streams
            if s.get("codec_type") == "video"
            and s.get("disposition", {}).get("attached_pic") != 1
        ),
        None,
    )
    if video is None:
        return None
    width, height, codec_name = video.get("width"), video.get("height"), video.get("codec_name")
    if not (isinstance(width, int) and isinstance(height, int) and codec_name):
        return None
    if _rotated_quarter_turn(video):
        width, height = height, width
    subtitles = tuple(
        SubtitleStream(index=i, codec_name=s.get("codec_name") or "")
        for i, s in enumerate(s for s in streams if s.get("codec_type") == "subtitle")
    )
    try:
        duration = float(data.get("format", {}).get("duration"))
    except (TypeError, ValueError):
        duration = None
    return MediaInfo(
        width=width,
        height=height,
        codec_name=codec_name,
        subtitles=subtitles,
        duration=duration,
        field_order=video.get("field_order"),
    )