ffmpeg_progress_yield

 19class FfmpegProgress:
 20    DUR_REGEX = re.compile(
 21        r"Duration: (?P<hour>\d{2}):(?P<min>\d{2}):(?P<sec>\d{2})\.(?P<ms>\d{2})"
 22    )
 23    TIME_REGEX = re.compile(
 24        r"out_time=(?P<hour>\d{2}):(?P<min>\d{2}):(?P<sec>\d{2})\.(?P<ms>\d{2})"
 25    )
 26    PROGRESS_REGEX = re.compile(r"[a-z0-9_]+=.+")
 27
 28    def __init__(
 29        self,
 30        cmd: List[str],
 31        dry_run: bool = False,
 32        exclude_progress: bool = False,
 33        ffprobe_path: str = "ffprobe",
 34    ) -> None:
 35        """Initialize the FfmpegProgress class.
 36
 37        Args:
 38            cmd (List[str]): A list of command line elements, e.g. ["ffmpeg", "-i", ...]
 39            dry_run (bool, optional): Only show what would be done. Defaults to False.
 40            exclude_progress (bool, optional): Exclude progress lines from output. Defaults to False.
 41            ffprobe_path (str, optional): Path to ffprobe executable. Defaults to "ffprobe".
 42        """
 43        self.cmd = cmd
 44        self.stderr: Union[str, None] = None
 45        self.dry_run = dry_run
 46        self.exclude_progress = exclude_progress
 47        self.ffprobe_path = ffprobe_path
 48        self.process: Any = None
 49        self.stderr_callback: Union[Callable[[str], None], None] = None
 50        self.base_popen_kwargs = {
 51            "stdin": subprocess.PIPE,  # Apply stdin isolation by creating separate pipe.
 52            "stdout": subprocess.PIPE,
 53            "stderr": subprocess.STDOUT,
 54            "universal_newlines": False,
 55        }
 56
 57        self.cmd_with_progress = (
 58            [self.cmd[0]] + ["-progress", "-", "-nostats"] + self.cmd[1:]
 59        )
 60        self.inputs_with_options = FfmpegProgress._get_inputs_with_options(self.cmd)
 61
 62        self.current_input_idx: int = 0
 63        self.total_dur: Union[None, int] = None
 64        # Skip probing duration in dry-run mode to avoid running ffprobe
 65        if not self.dry_run and FfmpegProgress._uses_error_loglevel(self.cmd):
 66            self.total_dur = self._probe_duration(self.cmd)
 67
 68        # Set up cleanup on garbage collection as a fallback
 69        self._cleanup_ref = weakref.finalize(self, self._cleanup_process, None)
 70
 71    @staticmethod
 72    def _cleanup_process(process: Any) -> None:
 73        """Clean up a process if it's still running."""
 74        if process is not None and hasattr(process, "poll"):
 75            try:
 76                if process.poll() is None:  # Process is still running
 77                    process.kill()
 78                    if hasattr(process, "wait"):
 79                        try:
 80                            process.wait(timeout=1.0)
 81                        except subprocess.TimeoutExpired:
 82                            pass  # Process didn't terminate gracefully, but we killed it
 83            except Exception:
 84                pass  # Ignore any errors during cleanup
 85
 86    def __del__(self) -> None:
 87        """Fallback cleanup when object is garbage collected."""
 88        if hasattr(self, "process") and self.process is not None:
 89            self._cleanup_process(self.process)
 90
 91    def __enter__(self) -> "FfmpegProgress":
 92        """Context manager entry."""
 93        return self
 94
 95    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
 96        """Context manager exit - ensures process cleanup."""
 97        if self.process is not None:
 98            try:
 99                if hasattr(self.process, "poll") and self.process.poll() is None:
100                    self.quit()
101            except Exception:
102                pass  # Ignore errors during cleanup
103
104    async def __aenter__(self) -> "FfmpegProgress":
105        """Async context manager entry."""
106        return self
107
108    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
109        """Async context manager exit - ensures process cleanup."""
110        if self.process is not None:
111            try:
112                if (
113                    hasattr(self.process, "returncode")
114                    and self.process.returncode is None
115                ):
116                    await self.async_quit()
117            except Exception:
118                pass  # Ignore errors during cleanup
119
120    def _process_output(
121        self,
122        stderr_line: str,
123        stderr: List[str],
124        duration_override: Union[float, None],
125    ) -> Union[float, None]:
126        """
127        Process the output of the ffmpeg command.
128
129        Args:
130            stderr_line (str): The line of stderr output.
131            stderr (List[str]): The list of stderr output.
132            duration_override (Union[float, None]): The duration of the video in seconds.
133
134        Returns:
135            Union[float, None]: The progress in percent.
136        """
137
138        if self.stderr_callback:
139            self.stderr_callback(stderr_line)
140
141        stderr.append(stderr_line.strip())
142        self.stderr = "\n".join(
143            filter(
144                lambda line: not (
145                    self.exclude_progress and self.PROGRESS_REGEX.match(line)
146                ),
147                stderr,
148            )
149        )
150
151        progress: Union[float, None] = None
152        # assign the total duration if it was found. this can happen multiple times for multiple inputs,
153        # in which case we have to determine the overall duration by taking the min/max (dependent on -shortest being present)
154        if (
155            current_dur_match := self.DUR_REGEX.search(stderr_line)
156        ) and duration_override is None:
157            input_options = self.inputs_with_options[self.current_input_idx]
158            current_dur_ms: int = to_ms(**current_dur_match.groupdict())
159            # if the previous line had "image2", it's a single image and we assume a really short intrinsic duration (4ms),
160            # but if it's a loop, we assume infinity
161            if "image2" in stderr[-2] and "-loop 1" in " ".join(input_options):
162                current_dur_ms = 2**64
163            if "-shortest" in self.cmd:
164                self.total_dur = (
165                    min(self.total_dur, current_dur_ms)
166                    if self.total_dur is not None
167                    else current_dur_ms
168                )
169            else:
170                self.total_dur = (
171                    max(self.total_dur, current_dur_ms)
172                    if self.total_dur is not None
173                    else current_dur_ms
174                )
175            self.current_input_idx += 1
176
177        if (
178            progress_time := self.TIME_REGEX.search(stderr_line)
179        ) and self.total_dur is not None:
180            elapsed_time = to_ms(**progress_time.groupdict())
181            progress = min(max(round(elapsed_time / self.total_dur * 100, 2), 0), 100)
182
183        return progress
184
185    def _probe_duration(self, cmd: List[str]) -> Optional[int]:
186        """
187        Get the duration via ffprobe from input media file
188        in case ffmpeg was run with loglevel=error.
189
190        Args:
191            cmd (List[str]): A list of command line elements, e.g. ["ffmpeg", "-i", ...]
192
193        Returns:
194            Optional[int]: The duration in milliseconds.
195        """
196        file_names = []
197        for i, arg in enumerate(cmd):
198            if arg == "-i":
199                file_name = cmd[i + 1]
200
201                # filter for filenames that we can probe, i.e. regular files
202                if os.path.isfile(file_name):
203                    file_names.append(file_name)
204
205        if len(file_names) == 0:
206            return None
207
208        durations = []
209
210        for file_name in file_names:
211            try:
212                output = subprocess.check_output(
213                    [
214                        self.ffprobe_path,
215                        "-loglevel",
216                        "error",
217                        "-hide_banner",
218                        "-show_entries",
219                        "format=duration",
220                        "-of",
221                        "default=noprint_wrappers=1:nokey=1",
222                        file_name,
223                    ],
224                    universal_newlines=True,
225                )
226                durations.append(int(float(output.strip()) * 1000))
227            except Exception:
228                # TODO: add logging
229                return None
230
231        return max(durations) if "-shortest" not in cmd else min(durations)
232
233    @staticmethod
234    def _uses_error_loglevel(cmd: List[str]) -> bool:
235        try:
236            idx = cmd.index("-loglevel")
237            if cmd[idx + 1] == "error":
238                return True
239            else:
240                return False
241        except ValueError:
242            return False
243
244    @staticmethod
245    def _get_inputs_with_options(cmd: List[str]) -> List[List[str]]:
246        """
247        Collect all inputs with their options.
248        For example, input is:
249
250            ffmpeg -i input1.mp4 -i input2.mp4 -i input3.mp4 -filter_complex ...
251
252        Output is:
253
254            [
255                ["-i", "input1.mp4"],
256                ["-i", "input2.mp4"],
257                ["-i", "input3.mp4"],
258            ]
259
260        Another example:
261
262            ffmpeg -f lavfi -i color=c=black:s=1920x1080 -loop 1 -i image.png -filter_complex ...
263
264        Output is:
265
266            [
267                ["-f", "lavfi", "-i", "color=c=black:s=1920x1080"],
268                ["-loop", "1", "-i", "image.png"],
269            ]
270        """
271        inputs = []
272        prev_index = 0
273        for i, arg in enumerate(cmd):
274            if arg == "-i":
275                inputs.append(cmd[prev_index : i + 2])
276                prev_index = i + 2
277
278        return inputs
279
280    def run_command_with_progress(
281        self, popen_kwargs=None, duration_override: Union[float, None] = None
282    ) -> Iterator[float]:
283        """
284        Run an ffmpeg command, trying to capture the process output and calculate
285        the duration / progress.
286        Yields the progress in percent.
287
288        Args:
289            popen_kwargs (dict, optional): A dict to specify extra arguments to the popen call, e.g. { creationflags: CREATE_NO_WINDOW }
290            duration_override (float, optional): The duration in seconds. If not specified, it will be calculated from the ffmpeg output.
291
292        Raises:
293            RuntimeError: If the command fails, an exception is raised.
294
295        Yields:
296            Iterator[float]: A generator that yields the progress in percent.
297        """
298        if self.dry_run:
299            yield from [0, 100]
300            return
301
302        if duration_override:
303            self.total_dur = int(duration_override * 1000)
304
305        base_popen_kwargs = self.base_popen_kwargs.copy()
306        if popen_kwargs is not None:
307            base_popen_kwargs.update(popen_kwargs)
308
309        self.process = subprocess.Popen(self.cmd_with_progress, **base_popen_kwargs)  # type: ignore
310
311        # Update the cleanup finalizer with the actual process
312        self._cleanup_ref.detach()
313        self._cleanup_ref = weakref.finalize(self, self._cleanup_process, self.process)
314
315        try:
316            yield 0
317
318            stderr: List[str] = []
319            while True:
320                if self.process.stdout is None:
321                    continue
322
323                stderr_line: str = (
324                    self.process.stdout.readline()
325                    .decode("utf-8", errors="replace")
326                    .strip()
327                )
328
329                if stderr_line == "" and self.process.poll() is not None:
330                    break
331
332                progress = self._process_output(stderr_line, stderr, duration_override)
333                if progress is not None:
334                    yield progress
335
336            if self.process.returncode != 0:
337                raise RuntimeError(f"Error running command {self.cmd}: {self.stderr}")
338
339            yield 100
340        finally:
341            # Ensure process cleanup even if an exception occurs
342            if self.process is not None:
343                try:
344                    if self.process.poll() is None:  # Process is still running
345                        self.process.kill()
346                        try:
347                            self.process.wait(timeout=1.0)
348                        except subprocess.TimeoutExpired:
349                            pass  # Process didn't terminate gracefully, but we killed it
350                except Exception:
351                    pass  # Ignore any errors during cleanup
352                finally:
353                    self.process = None
354                    # Detach the finalizer since we've cleaned up manually
355                    if hasattr(self, "_cleanup_ref"):
356                        self._cleanup_ref.detach()
357
358    async def async_run_command_with_progress(
359        self, popen_kwargs=None, duration_override: Union[float, None] = None
360    ) -> AsyncIterator[float]:
361        """
362        Asynchronously run an ffmpeg command, trying to capture the process output and calculate
363        the duration / progress.
364        Yields the progress in percent.
365
366        Args:
367            popen_kwargs (dict, optional): A dict to specify extra arguments to the popen call, e.g. { creationflags: CREATE_NO_WINDOW }
368            duration_override (float, optional): The duration in seconds. If not specified, it will be calculated from the ffmpeg output.
369
370        Raises:
371            RuntimeError: If the command fails, an exception is raised.
372        """
373        if self.dry_run:
374            yield 0
375            yield 100
376            return
377
378        if duration_override:
379            self.total_dur = int(duration_override * 1000)
380
381        base_popen_kwargs = self.base_popen_kwargs.copy()
382        if popen_kwargs is not None:
383            base_popen_kwargs.update(popen_kwargs)
384
385        # Remove stdout and stderr from base_popen_kwargs as we're setting them explicitly
386        base_popen_kwargs.pop("stdout", None)
387        base_popen_kwargs.pop("stderr", None)
388
389        self.process = await asyncio.create_subprocess_exec(
390            *self.cmd_with_progress,
391            stdout=asyncio.subprocess.PIPE,
392            stderr=asyncio.subprocess.STDOUT,
393            **base_popen_kwargs,  # type: ignore
394        )
395
396        # Update the cleanup finalizer with the actual process
397        self._cleanup_ref.detach()
398        self._cleanup_ref = weakref.finalize(self, self._cleanup_process, self.process)
399
400        try:
401            yield 0
402
403            stderr: List[str] = []
404            while True:
405                if self.process.stdout is None:
406                    continue
407
408                stderr_line: Union[bytes, None] = await self.process.stdout.readline()
409                if not stderr_line:
410                    # Process has finished, check the return code
411                    await self.process.wait()
412                    if self.process.returncode != 0:
413                        raise RuntimeError(
414                            f"Error running command {self.cmd}: {self.stderr}"
415                        )
416                    break
417                stderr_line_str = stderr_line.decode("utf-8", errors="replace").strip()
418
419                progress = self._process_output(
420                    stderr_line_str, stderr, duration_override
421                )
422                if progress is not None:
423                    yield progress
424
425            yield 100
426        except GeneratorExit:
427            # Handle case where async generator is closed prematurely
428            await self._async_cleanup_process()
429            raise
430        except Exception:
431            # Handle any other exception
432            await self._async_cleanup_process()
433            raise
434        finally:
435            # Normal cleanup
436            await self._async_cleanup_process()
437
438    async def _async_cleanup_process(self) -> None:
439        """Clean up the async process."""
440        if self.process is not None:
441            try:
442                if self.process.returncode is None:  # Process is still running
443                    self.process.kill()
444                    try:
445                        await self.process.wait()
446                    except Exception:
447                        pass  # Ignore any errors during cleanup
448            except Exception:
449                pass  # Ignore any errors during cleanup
450            finally:
451                self.process = None
452                # Detach the finalizer since we've cleaned up manually
453                if hasattr(self, "_cleanup_ref"):
454                    self._cleanup_ref.detach()
455
456    def quit_gracefully(self) -> None:
457        """
458        Quit the ffmpeg process by sending 'q'
459
460        Raises:
461            RuntimeError: If no process is found.
462        """
463        if self.process is None:
464            raise RuntimeError("No process found. Did you run the command?")
465
466        self.process.communicate(input=b"q")
467        self.process.kill()
468        self.process = None
469
470    def quit(self) -> None:
471        """
472        Quit the ffmpeg process by sending SIGKILL.
473
474        Raises:
475            RuntimeError: If no process is found.
476        """
477        if self.process is None:
478            raise RuntimeError("No process found. Did you run the command?")
479
480        self.process.kill()
481        self.process = None
482
483    async def async_quit_gracefully(self) -> None:
484        """
485        Quit the ffmpeg process by sending 'q' asynchronously
486
487        Raises:
488            RuntimeError: If no process is found.
489        """
490        if self.process is None:
491            raise RuntimeError("No process found. Did you run the command?")
492
493        self.process.stdin.write(b"q")
494        await self.process.stdin.drain()
495        await self.process.wait()
496        self.process = None
497
498    async def async_quit(self) -> None:
499        """
500        Quit the ffmpeg process by sending SIGKILL asynchronously.
501
502        Raises:
503            RuntimeError: If no process is found.
504        """
505        if self.process is None:
506            raise RuntimeError("No process found. Did you run the command?")
507
508        self.process.kill()
509        await self.process.wait()
510        self.process = None
511
512    def set_stderr_callback(self, callback: Callable[[str], None]) -> None:
513        """
514        Set a callback function to be called on stderr output.
515        The callback function must accept a single string argument.
516        Note that this is called on every line of stderr output, so it can be called a lot.
517        Also note that stdout/stderr are joined into one stream, so you might get stdout output in the callback.
518
519        Args:
520            callback (Callable[[str], None]): A callback function that accepts a single string argument.
521        """
522        if not callable(callback) or len(callback.__code__.co_varnames) != 1:
523            raise ValueError(
524                "Callback must be a function that accepts only one argument"
525            )
526
527        self.stderr_callback = callback