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