185 lines
6.1 KiB
Python
185 lines
6.1 KiB
Python
"""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
|