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