Version 1 OneShot

ffwrap/probe.py

@@ -0,0 +1,184 @@
+"""ffprobe wrapper producing simple dataclasses."""
+
+from __future__ import annotations
+
+import json
+import subprocess
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from .formats import describe_channels
+from .tools import run_quiet
+
+
+class ProbeError(Exception):
+    pass
+
+
+@dataclass
+class Stream:
+    index: int                 # absolute stream index inside the file
+    kind: str                  # "video" | "audio" | "subtitle" | "data" | "attachment"
+    codec: str | None
+    profile: str | None = None
+    language: str | None = None
+    title: str | None = None
+    default: bool = False
+    forced: bool = False
+    attached_pic: bool = False
+    # video
+    width: int | None = None
+    height: int | None = None
+    fps: float | None = None
+    rotation: int = 0
+    pix_fmt: str | None = None
+    # audio
+    channels: int | None = None
+    channel_layout: str | None = None
+    sample_rate: int | None = None
+    bit_rate: int | None = None
+
+    @property
+    def display_size(self) -> tuple[int, int] | None:
+        """Width/height after applying rotation metadata (what players show)."""
+        if not self.width or not self.height:
+            return None
+        if self.rotation % 180 == 90:
+            return self.height, self.width
+        return self.width, self.height
+
+    def describe(self) -> str:
+        if self.kind == "audio":
+            parts = [(self.codec or "?").upper(), describe_channels(self.channels, self.channel_layout)]
+            if self.sample_rate:
+                parts.append(f"{self.sample_rate / 1000:g} kHz")
+            if self.bit_rate:
+                parts.append(f"{self.bit_rate // 1000} kb/s")
+            return " ".join(parts)
+        if self.kind == "video":
+            size = self.display_size
+            s = (self.codec or "?").upper()
+            if size:
+                s += f" {size[0]}x{size[1]}"
+            if self.fps:
+                s += f" {self.fps:.3g}fps"
+            return s
+        return self.codec or "?"
+
+
+@dataclass
+class MediaInfo:
+    path: Path
+    duration: float | None
+    format_name: str
+    streams: list[Stream] = field(default_factory=list)
+
+    @property
+    def video(self) -> Stream | None:
+        """The main video stream (cover art / thumbnails are ignored)."""
+        for s in self.streams:
+            if s.kind == "video" and not s.attached_pic:
+                return s
+        return None
+
+    @property
+    def audio(self) -> list[Stream]:
+        return [s for s in self.streams if s.kind == "audio"]
+
+    @property
+    def subtitles(self) -> list[Stream]:
+        return [s for s in self.streams if s.kind == "subtitle"]
+
+    def stream(self, index: int) -> Stream | None:
+        for s in self.streams:
+            if s.index == index:
+                return s
+        return None
+
+
+def _to_float(v) -> float | None:
+    try:
+        f = float(v)
+    except (TypeError, ValueError):
+        return None
+    return f if f == f else None   # filter NaN
+
+
+def _to_int(v) -> int | None:
+    f = _to_float(v)
+    return int(f) if f is not None else None
+
+
+def _rate(v: str | None) -> float | None:
+    if not v or v in ("0/0", "0"):
+        return None
+    if "/" in v:
+        num, den = v.split("/", 1)
+        n, d = _to_float(num), _to_float(den)
+        return n / d if n and d else None
+    return _to_float(v)
+
+
+def _rotation(raw: dict) -> int:
+    for sd in raw.get("side_data_list") or []:
+        if "rotation" in sd:
+            return int(_to_float(sd["rotation"]) or 0) % 360
+    rot = (raw.get("tags") or {}).get("rotate")
+    return int(_to_float(rot) or 0) % 360
+
+
+def parse_probe(path: Path, data: dict) -> MediaInfo:
+    fmt = data.get("format") or {}
+    streams: list[Stream] = []
+    for raw in data.get("streams") or []:
+        tags = {k.lower(): v for k, v in (raw.get("tags") or {}).items()}
+        disp = raw.get("disposition") or {}
+        lang = tags.get("language")
+        s = Stream(
+            index=int(raw.get("index", len(streams))),
+            kind=raw.get("codec_type") or "data",
+            codec=raw.get("codec_name"),
+            profile=raw.get("profile"),
+            language=lang if lang and lang != "und" else None,
+            title=tags.get("title"),
+            default=bool(disp.get("default")),
+            forced=bool(disp.get("forced")),
+            attached_pic=bool(disp.get("attached_pic")),
+            width=_to_int(raw.get("width")),
+            height=_to_int(raw.get("height")),
+            fps=_rate(raw.get("avg_frame_rate")) or _rate(raw.get("r_frame_rate")),
+            rotation=_rotation(raw),
+            pix_fmt=raw.get("pix_fmt"),
+            channels=_to_int(raw.get("channels")),
+            channel_layout=raw.get("channel_layout"),
+            sample_rate=_to_int(raw.get("sample_rate")),
+            bit_rate=_to_int(raw.get("bit_rate")) or _to_int(tags.get("bps")),
+        )
+        streams.append(s)
+
+    duration = _to_float(fmt.get("duration"))
+    if not duration:
+        durations = [_to_float(r.get("duration")) for r in data.get("streams") or []]
+        duration = max((d for d in durations if d), default=None)
+    return MediaInfo(path=path, duration=duration, format_name=fmt.get("format_name", ""), streams=streams)
+
+
+def probe(ffprobe: str, path: Path) -> MediaInfo:
+    args = [ffprobe, "-v", "error", "-print_format", "json", "-show_format", "-show_streams", str(path)]
+    try:
+        res = run_quiet(args, timeout=120)
+    except (OSError, subprocess.SubprocessError) as e:
+        raise ProbeError(str(e)) from e
+    if res.returncode != 0:
+        lines = res.stderr.decode(errors="replace").strip().splitlines()
+        msg = lines[-1] if lines else f"ffprobe exited with code {res.returncode}"
+        # ffprobe prefixes errors with the file name, which the queue already shows.
+        raise ProbeError(msg.removeprefix(f"{path}: "))
+    try:
+        data = json.loads(res.stdout.decode("utf-8", errors="replace"))
+    except json.JSONDecodeError as e:
+        raise ProbeError(f"Could not parse ffprobe output: {e}") from e
+    info = parse_probe(path, data)
+    if not info.video and not info.audio:
+        raise ProbeError("No video or audio streams found")
+    return info