From 2f3331c6194a9d977e2e06541320830feb9ac7ad Mon Sep 17 00:00:00 2001 From: Tao Peng Date: Tue, 11 Oct 2022 15:16:29 -0700 Subject: [PATCH] feat: add option to allow capture stderr from ffmpeg/ffprobe (#565) * feat: add option to allow capture stderr from ffmpeg/ffprobe * add tests * fix tests * add comments * fix decoding error * fix * fix --- mapillary_tools/ffmpeg.py | 105 +++++++++++++++++++++++++++----------- tests/unit/test_ffmpeg.py | 75 +++++++++++++++++++++++++++ 2 files changed, 151 insertions(+), 29 deletions(-) create mode 100644 tests/unit/test_ffmpeg.py diff --git a/mapillary_tools/ffmpeg.py b/mapillary_tools/ffmpeg.py index 00245509..a6c428c6 100644 --- a/mapillary_tools/ffmpeg.py +++ b/mapillary_tools/ffmpeg.py @@ -34,52 +34,117 @@ class Stream(TypedDict): width: int -class ProbeFormat(TypedDict): - filename: str - duration: str - - class ProbeOutput(TypedDict): streams: T.List[Stream] - format: T.Dict class FFmpegNotFoundError(Exception): pass +_MAX_STDERR_LENGTH = 2048 + + +def _truncate_begin(s: str) -> str: + if _MAX_STDERR_LENGTH < len(s): + return "..." + s[-_MAX_STDERR_LENGTH:] + else: + return s + + +def _truncate_end(s: str) -> str: + if _MAX_STDERR_LENGTH < len(s): + return s[:_MAX_STDERR_LENGTH] + "..." + else: + return s + + +class FFmpegCalledProcessError(Exception): + def __init__(self, ex: subprocess.CalledProcessError): + self.inner_ex = ex + + def __str__(self) -> str: + msg = str(self.inner_ex) + if self.inner_ex.stderr is not None: + try: + stderr = self.inner_ex.stderr.decode("utf-8") + except UnicodeDecodeError: + stderr = str(self.inner_ex.stderr) + msg += f"\nSTDERR: {_truncate_begin(stderr)}" + return msg + + class FFMPEG: def __init__( - self, ffmpeg_path: str = "ffmpeg", ffprobe_path: str = "ffprobe" + self, + ffmpeg_path: str = "ffmpeg", + ffprobe_path: str = "ffprobe", + stderr: T.Optional[int] = None, ) -> None: + """ + ffmpeg_path: path to ffmpeg binary + ffprobe_path: path to ffprobe binary + stderr: param passed to subprocess.run to control whether to capture stderr + """ self.ffmpeg_path = ffmpeg_path self.ffprobe_path = ffprobe_path + self.stderr = stderr def _run_ffprobe_json(self, cmd: T.List[str]) -> T.Dict: full_cmd = [self.ffprobe_path, "-print_format", "json", *cmd] LOG.info(f"Extracting video information: {' '.join(full_cmd)}") try: - output = subprocess.check_output(full_cmd) + completed = subprocess.run( + full_cmd, + check=True, + stdout=subprocess.PIPE, + stderr=self.stderr, + ) except FileNotFoundError: raise FFmpegNotFoundError( f'The ffprobe command "{self.ffprobe_path}" not found' ) + except subprocess.CalledProcessError as ex: + raise FFmpegCalledProcessError(ex) from ex + try: - return json.loads(output) + stdout = completed.stdout.decode("utf-8") + except UnicodeDecodeError: + raise RuntimeError( + f"Error decoding ffprobe output as unicode: {_truncate_end(str(completed.stdout))}" + ) + + try: + output = json.loads(stdout) except json.JSONDecodeError: raise RuntimeError( - f"Error JSON decoding ffprobe output: {output.decode('utf-8')}" + f"Error JSON decoding ffprobe output: {_truncate_end(stdout)}" + ) + + # This check is for macOS: + # ffprobe -hide_banner -print_format json not_exists + # you will get exit code == 0 with the following stdout and stderr: + # { + # } + # not_exists: No such file or directory + if not output: + raise RuntimeError( + f"Empty JSON ffprobe output with STDERR: {_truncate_begin(str(completed.stderr))}" ) + return output + def _run_ffmpeg(self, cmd: T.List[str]) -> None: full_cmd = [self.ffmpeg_path, *cmd] LOG.info(f"Extracting frames: {' '.join(full_cmd)}") try: - subprocess.check_call(full_cmd) + subprocess.run(full_cmd, check=True, stderr=self.stderr) except FileNotFoundError: raise FFmpegNotFoundError( f'The ffmpeg command "{self.ffmpeg_path}" not found' ) + except subprocess.CalledProcessError as ex: + raise FFmpegCalledProcessError(ex) from ex def probe_format_and_streams(self, video_path: Path) -> ProbeOutput: cmd = [ @@ -90,24 +155,6 @@ def probe_format_and_streams(self, video_path: Path) -> ProbeOutput: ] return T.cast(ProbeOutput, self._run_ffprobe_json(cmd)) - def extract_stream(self, source: Path, dest: Path, stream_id: int) -> None: - cmd = [ - "-hide_banner", - "-i", - str(source), - "-y", # overwrite - potentially dangerous - "-nostats", - "-codec", - "copy", - "-map", - f"0:{stream_id}", - "-f", - "rawvideo", - str(dest), - ] - - self._run_ffmpeg(cmd) - def extract_frames( self, video_path: Path, diff --git a/tests/unit/test_ffmpeg.py b/tests/unit/test_ffmpeg.py new file mode 100644 index 00000000..036e97ae --- /dev/null +++ b/tests/unit/test_ffmpeg.py @@ -0,0 +1,75 @@ +import os +import subprocess +from pathlib import Path + +import pytest + +from mapillary_tools import ffmpeg + + +def ffmpeg_installed(): + ffmpeg_path = os.getenv("MAPILLARY_TOOLS_FFMPEG_PATH", "ffmpeg") + ffprobe_path = os.getenv("MAPILLARY_TOOLS_FFPROBE_PATH", "ffprobe") + try: + subprocess.run([ffmpeg_path, "-version"]) + # In Windows, ffmpeg is installed but ffprobe is not? + subprocess.run([ffprobe_path, "-version"]) + except FileNotFoundError: + return False + return True + + +is_ffmpeg_installed = ffmpeg_installed() + + +def test_ffmpeg_not_exists(): + if not is_ffmpeg_installed: + pytest.skip("ffmpeg not installed") + + ff = ffmpeg.FFMPEG() + try: + ff.extract_frames(Path("not_exist_a"), Path("not_exist_b"), sample_interval=2) + except ffmpeg.FFmpegCalledProcessError as ex: + assert "STDERR:" not in str(ex) + else: + assert False, "FFmpegCalledProcessError not raised" + + ff = ffmpeg.FFMPEG(stderr=subprocess.PIPE) + try: + ff.extract_frames(Path("not_exist_a"), Path("not_exist_b"), sample_interval=2) + except ffmpeg.FFmpegCalledProcessError as ex: + assert "STDERR:" in str(ex) + else: + assert False, "FFmpegCalledProcessError not raised" + + +def test_ffprobe_not_exists(): + if not is_ffmpeg_installed: + pytest.skip("ffmpeg not installed") + + ff = ffmpeg.FFMPEG() + try: + x = ff.probe_format_and_streams(Path("not_exist_a")) + except ffmpeg.FFmpegCalledProcessError as ex: + # exc from linux + assert "STDERR:" not in str(ex) + except RuntimeError as ex: + # exc from macos + assert "Empty JSON ffprobe output with STDERR: None" == str(ex) + else: + assert False, "RuntimeError not raised" + + ff = ffmpeg.FFMPEG(stderr=subprocess.PIPE) + try: + x = ff.probe_format_and_streams(Path("not_exist_a")) + except ffmpeg.FFmpegCalledProcessError as ex: + # exc from linux + assert "STDERR:" in str(ex) + except RuntimeError as ex: + # exc from macos + assert ( + "Empty JSON ffprobe output with STDERR: b'not_exist_a: No such file or directory" + in str(ex) + ) + else: + assert False, "RuntimeError not raised"