ffmpeg_progress_yield

 18class FfmpegProgress:
 19    DUR_REGEX = re.compile(
 20        r"Duration: (?P<hour>\d{2}):(?P<min>\d{2}):(?P<sec>\d{2})\.(?P<ms>\d{2})"
 21    )
 22    TIME_REGEX = re.compile(
 23        r"out_time=(?P<hour>\d{2}):(?P<min>\d{2}):(?P<sec>\d{2})\.(?P<ms>\d{2})"
 24    )
 25    PROGRESS_REGEX = re.compile(
 26        r"[a-z0-9_]+=.+"
 27    )
 28
 29    def __init__(self, cmd: List[str], dry_run: bool = False, exclude_progress: bool = False) -> None:
 30        """Initialize the FfmpegProgress class.
 31
 32        Args:
 33            cmd (List[str]): A list of command line elements, e.g. ["ffmpeg", "-i", ...]
 34            dry_run (bool, optional): Only show what would be done. Defaults to False.
 35        """
 36        self.cmd = cmd
 37        self.stderr: Union[str, None] = None
 38        self.dry_run = dry_run
 39        self.exclude_progress = exclude_progress
 40        self.process: Any = None
 41        self.stderr_callback: Union[Callable[[str], None], None] = None
 42        self.base_popen_kwargs = {
 43            "stdin": subprocess.PIPE,  # Apply stdin isolation by creating separate pipe.
 44            "stdout": subprocess.PIPE,
 45            "stderr": subprocess.STDOUT,
 46            "universal_newlines": False,
 47        }
 48
 49        self.cmd_with_progress = (
 50            [self.cmd[0]] + ["-progress", "-", "-nostats"] + self.cmd[1:]
 51        )
 52        self.inputs_with_options = FfmpegProgress._get_inputs_with_options(self.cmd)
 53
 54        self.current_input_idx: int = 0
 55        self.total_dur: Union[None, int] = None
 56        if FfmpegProgress._uses_error_loglevel(self.cmd):
 57            self.total_dur = FfmpegProgress._probe_duration(self.cmd)
 58
 59    def _process_output(
 60        self,
 61        stderr_line: str,
 62        stderr: List[str],
 63        duration_override: Union[float, None],
 64    ) -> Union[float, None]:
 65        """
 66        Process the output of the ffmpeg command.
 67
 68        Args:
 69            stderr_line (str): The line of stderr output.
 70            stderr (List[str]): The list of stderr output.
 71            duration_override (Union[float, None]): The duration of the video in seconds.
 72
 73        Returns:
 74            Union[float, None]: The progress in percent.
 75        """
 76
 77        if self.stderr_callback:
 78            self.stderr_callback(stderr_line)
 79
 80        stderr.append(stderr_line.strip())
 81        self.stderr = "\n".join(
 82            filter(
 83                lambda line: not (self.exclude_progress and self.PROGRESS_REGEX.match(line)),
 84                stderr
 85            )
 86        )
 87
 88        progress: Union[float, None] = None
 89        # assign the total duration if it was found. this can happen multiple times for multiple inputs,
 90        # in which case we have to determine the overall duration by taking the min/max (dependent on -shortest being present)
 91        if (
 92            current_dur_match := self.DUR_REGEX.search(stderr_line)
 93        ) and duration_override is None:
 94            input_options = self.inputs_with_options[self.current_input_idx]
 95            current_dur_ms: int = to_ms(**current_dur_match.groupdict())
 96            # if the previous line had "image2", it's a single image and we assume a really short intrinsic duration (4ms),
 97            # but if it's a loop, we assume infinity
 98            if "image2" in stderr[-2] and "-loop 1" in " ".join(input_options):
 99                current_dur_ms = 2**64
100            if "-shortest" in self.cmd:
101                self.total_dur = (
102                    min(self.total_dur, current_dur_ms)
103                    if self.total_dur is not None
104                    else current_dur_ms
105                )
106            else:
107                self.total_dur = (
108                    max(self.total_dur, current_dur_ms)
109                    if self.total_dur is not None
110                    else current_dur_ms
111                )
112            self.current_input_idx += 1
113
114        if (
115            progress_time := self.TIME_REGEX.search(stderr_line)
116        ) and self.total_dur is not None:
117            elapsed_time = to_ms(**progress_time.groupdict())
118            progress = min(max(round(elapsed_time / self.total_dur * 100, 2), 0), 100)
119
120        return progress
121
122    @staticmethod
123    def _probe_duration(cmd: List[str]) -> Optional[int]:
124        """
125        Get the duration via ffprobe from input media file
126        in case ffmpeg was run with loglevel=error.
127
128        Args:
129            cmd (List[str]): A list of command line elements, e.g. ["ffmpeg", "-i", ...]
130
131        Returns:
132            Optional[int]: The duration in milliseconds.
133        """
134        file_names = []
135        for i, arg in enumerate(cmd):
136            if arg == "-i":
137                file_name = cmd[i + 1]
138
139                # filter for filenames that we can probe, i.e. regular files
140                if os.path.isfile(file_name):
141                    file_names.append(file_name)
142
143        if len(file_names) == 0:
144            return None
145
146        durations = []
147
148        for file_name in file_names:
149            try:
150                output = subprocess.check_output(
151                    [
152                        "ffprobe",
153                        "-loglevel",
154                        "error",
155                        "-hide_banner",
156                        "-show_entries",
157                        "format=duration",
158                        "-of",
159                        "default=noprint_wrappers=1:nokey=1",
160                        file_name,
161                    ],
162                    universal_newlines=True,
163                )
164                durations.append(int(float(output.strip()) * 1000))
165            except Exception:
166                # TODO: add logging
167                return None
168
169        return max(durations) if "-shortest" not in cmd else min(durations)
170
171    @staticmethod
172    def _uses_error_loglevel(cmd: List[str]) -> bool:
173        try:
174            idx = cmd.index("-loglevel")
175            if cmd[idx + 1] == "error":
176                return True
177            else:
178                return False
179        except ValueError:
180            return False
181
182    @staticmethod
183    def _get_inputs_with_options(cmd: List[str]) -> List[List[str]]:
184        """
185        Collect all inputs with their options.
186        For example, input is:
187
188            ffmpeg -i input1.mp4 -i input2.mp4 -i input3.mp4 -filter_complex ...
189
190        Output is:
191
192            [
193                ["-i", "input1.mp4"],
194                ["-i", "input2.mp4"],
195                ["-i", "input3.mp4"],
196            ]
197
198        Another example:
199
200            ffmpeg -f lavfi -i color=c=black:s=1920x1080 -loop 1 -i image.png -filter_complex ...
201
202        Output is:
203
204            [
205                ["-f", "lavfi", "-i", "color=c=black:s=1920x1080"],
206                ["-loop", "1", "-i", "image.png"],
207            ]
208        """
209        inputs = []
210        prev_index = 0
211        for i, arg in enumerate(cmd):
212            if arg == "-i":
213                inputs.append(cmd[prev_index : i + 2])
214                prev_index = i + 2
215
216        return inputs
217
218    def run_command_with_progress(
219        self, popen_kwargs=None, duration_override: Union[float, None] = None
220    ) -> Iterator[float]:
221        """
222        Run an ffmpeg command, trying to capture the process output and calculate
223        the duration / progress.
224        Yields the progress in percent.
225
226        Args:
227            popen_kwargs (dict, optional): A dict to specify extra arguments to the popen call, e.g. { creationflags: CREATE_NO_WINDOW }
228            duration_override (float, optional): The duration in seconds. If not specified, it will be calculated from the ffmpeg output.
229
230        Raises:
231            RuntimeError: If the command fails, an exception is raised.
232
233        Yields:
234            Iterator[float]: A generator that yields the progress in percent.
235        """
236        if self.dry_run:
237            yield from [0, 100]
238            return
239
240        if duration_override:
241            self.total_dur = int(duration_override * 1000)
242
243        base_popen_kwargs = self.base_popen_kwargs.copy()
244        if popen_kwargs is not None:
245            base_popen_kwargs.update(popen_kwargs)
246
247        self.process = subprocess.Popen(self.cmd_with_progress, **base_popen_kwargs)  # type: ignore
248
249        yield 0
250
251        stderr: List[str] = []
252        while True:
253            if self.process.stdout is None:
254                continue
255
256            stderr_line: str = (
257                self.process.stdout.readline().decode("utf-8", errors="replace").strip()
258            )
259
260            if stderr_line == "" and self.process.poll() is not None:
261                break
262
263            progress = self._process_output(stderr_line, stderr, duration_override)
264            if progress is not None:
265                yield progress
266
267        if self.process.returncode != 0:
268            raise RuntimeError(f"Error running command {self.cmd}: {self.stderr}")
269
270        yield 100
271        self.process = None
272
273    async def async_run_command_with_progress(
274        self, popen_kwargs=None, duration_override: Union[float, None] = None
275    ) -> AsyncIterator[float]:
276        """
277        Asynchronously 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        if self.dry_run:
289            yield 0
290            yield 100
291            return
292
293        if duration_override:
294            self.total_dur = int(duration_override * 1000)
295
296        base_popen_kwargs = self.base_popen_kwargs.copy()
297        if popen_kwargs is not None:
298            base_popen_kwargs.update(popen_kwargs)
299
300        # Remove stdout and stderr from base_popen_kwargs as we're setting them explicitly
301        base_popen_kwargs.pop("stdout", None)
302        base_popen_kwargs.pop("stderr", None)
303
304        self.process = await asyncio.create_subprocess_exec(
305            *self.cmd_with_progress,
306            stdout=asyncio.subprocess.PIPE,
307            stderr=asyncio.subprocess.STDOUT,
308            **base_popen_kwargs,  # type: ignore
309        )
310
311        yield 0
312
313        stderr: List[str] = []
314        while True:
315            if self.process.stdout is None:
316                continue
317
318            stderr_line: Union[bytes, None] = await self.process.stdout.readline()
319            if not stderr_line:
320                # Process has finished, check the return code
321                await self.process.wait()
322                if self.process.returncode != 0:
323                    raise RuntimeError(
324                        f"Error running command {self.cmd}: {self.stderr}"
325                    )
326                break
327            stderr_line_str = stderr_line.decode("utf-8", errors="replace").strip()
328
329            progress = self._process_output(stderr_line_str, stderr, duration_override)
330            if progress is not None:
331                yield progress
332
333        yield 100
334        self.process = None
335
336    def quit_gracefully(self) -> None:
337        """
338        Quit the ffmpeg process by sending 'q'
339
340        Raises:
341            RuntimeError: If no process is found.
342        """
343        if self.process is None:
344            raise RuntimeError("No process found. Did you run the command?")
345
346        self.process.communicate(input=b"q")
347        self.process.kill()
348        self.process = None
349
350    def quit(self) -> None:
351        """
352        Quit the ffmpeg process by sending SIGKILL.
353
354        Raises:
355            RuntimeError: If no process is found.
356        """
357        if self.process is None:
358            raise RuntimeError("No process found. Did you run the command?")
359
360        self.process.kill()
361        self.process = None
362
363    async def async_quit_gracefully(self) -> None:
364        """
365        Quit the ffmpeg process by sending 'q' asynchronously
366
367        Raises:
368            RuntimeError: If no process is found.
369        """
370        if self.process is None:
371            raise RuntimeError("No process found. Did you run the command?")
372
373        self.process.stdin.write(b"q")
374        await self.process.stdin.drain()
375        await self.process.wait()
376        self.process = None
377
378    async def async_quit(self) -> None:
379        """
380        Quit the ffmpeg process by sending SIGKILL asynchronously.
381
382        Raises:
383            RuntimeError: If no process is found.
384        """
385        if self.process is None:
386            raise RuntimeError("No process found. Did you run the command?")
387
388        self.process.kill()
389        await self.process.wait()
390        self.process = None
391
392    def set_stderr_callback(self, callback: Callable[[str], None]) -> None:
393        """
394        Set a callback function to be called on stderr output.
395        The callback function must accept a single string argument.
396        Note that this is called on every line of stderr output, so it can be called a lot.
397        Also note that stdout/stderr are joined into one stream, so you might get stdout output in the callback.
398
399        Args:
400            callback (Callable[[str], None]): A callback function that accepts a single string argument.
401        """
402        if not callable(callback) or len(callback.__code__.co_varnames) != 1:
403            raise ValueError(
404                "Callback must be a function that accepts only one argument"
405            )
406
407        self.stderr_callback = callback