diff --git a/mediapy/__init__.py b/mediapy/__init__.py index af22117..266ad2d 100644 --- a/mediapy/__init__.py +++ b/mediapy/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 The mediapy Authors. +# Copyright 2026 The mediapy Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -1192,6 +1192,7 @@ def _run_ffmpeg( encoding: None = None, # No encoding -> bytes allowed_input_files: Sequence[str] | None = None, allowed_output_files: Sequence[str] | None = None, + sandbox_max_run_time_secs: int | None = None, ) -> subprocess.Popen[bytes]: ... @@ -1205,6 +1206,7 @@ def _run_ffmpeg( encoding: str = ..., # Encoding -> str allowed_input_files: Sequence[str] | None = None, allowed_output_files: Sequence[str] | None = None, + sandbox_max_run_time_secs: int | None = None, ) -> subprocess.Popen[str]: ... @@ -1217,6 +1219,7 @@ def _run_ffmpeg( encoding: str | None = None, allowed_input_files: Sequence[str] | None = None, allowed_output_files: Sequence[str] | None = None, + sandbox_max_run_time_secs: int | None = None, ) -> subprocess.Popen[bytes] | subprocess.Popen[str]: """Runs ffmpeg with the given args. @@ -1228,6 +1231,8 @@ def _run_ffmpeg( encoding: Same as in `subprocess.Popen`. allowed_input_files: The input files to allow for ffmpeg. allowed_output_files: The output files to allow for ffmpeg. + sandbox_max_run_time_secs: The maximum time in seconds to run the sandbox. + If None, the default limit is 30 minutes. Returns: The subprocess.Popen object with running ffmpeg process. @@ -1238,9 +1243,11 @@ def _run_ffmpeg( env: Any = None # pylint: disable=unused-variable ffmpeg_path = _get_ffmpeg_path() - # Allowed input and output files are not supported in open source. + # Sandbox max runtime, allowed input and ouput files are not supported in + # open source. del allowed_input_files del allowed_output_files + del sandbox_max_run_time_secs argv.append(ffmpeg_path) argv.extend(ffmpeg_args) @@ -1402,6 +1409,8 @@ class VideoReader(_VideoIO): bps: The estimated bitrate of the video stream in bits per second, retrieved from the video header. stream_index: The stream index to read from. The default is 0. + sandbox_max_run_time_secs: The maximum time in seconds to run the sandbox. + If None, the default limit is 30 minutes. Unused in open source. """ path_or_url: _Path @@ -1422,6 +1431,7 @@ def __init__( stream_index: int = 0, output_format: str = 'rgb', dtype: _DTypeLike = np.uint8, + sandbox_max_run_time_secs: int | None = None, ): if output_format not in {'rgb', 'yuv', 'gray'}: raise ValueError( @@ -1433,6 +1443,7 @@ def __init__( self.dtype = np.dtype(dtype) if self.dtype.type not in (np.uint8, np.uint16): raise ValueError(f'Type {dtype} is not np.uint8 or np.uint16.') + self.sandbox_max_run_time_secs = sandbox_max_run_time_secs self._read_via_local_file: Any = None self._popen: subprocess.Popen[bytes] | None = None self._proc: subprocess.Popen[bytes] | None = None @@ -1475,6 +1486,7 @@ def __enter__(self) -> 'VideoReader': stdout=subprocess.PIPE, stderr=subprocess.PIPE, allowed_input_files=[tmp_name], + sandbox_max_run_time_secs=self.sandbox_max_run_time_secs, ) self._proc = self._popen.__enter__() except Exception: @@ -1491,6 +1503,9 @@ def read(self) -> _NDArray | None: Returns: A numpy array in the format specified by `output_format`, i.e., a 3D array with 3 color channels, except for format 'gray' which is 2D. + + Raises: + RuntimeError: If there is an error reading from the output file. """ assert self._proc, 'Error: reading from an already closed context.' stdout = self._proc.stdout @@ -1499,7 +1514,17 @@ def read(self) -> _NDArray | None: if not data: # Due to either end-of-file or subprocess error. self.close() # Raises exception if subprocess had error. return None # To indicate end-of-file. - assert len(data) == self._num_bytes_per_image + if len(data) != self._num_bytes_per_image: + self._proc.wait() + stderr = self._proc.stderr + stderr_output = '' + if stderr is not None: + stderr_output = stderr.read().decode('utf-8', errors='replace').strip() + raise RuntimeError( + f'ffmpeg exited with code {self._proc.returncode}.\nIncomplete' + f' frame read: expected {self._num_bytes_per_image} bytes, but got' + f' {len(data)}.\nffmpeg stderr:\n{stderr_output}' + ) image = np.frombuffer(data, dtype=self.dtype) if self.output_format == 'rgb': image = image.reshape(*self.shape, 3) @@ -1574,6 +1599,8 @@ class VideoWriter(_VideoIO): 'yuv420p' (2x2-subsampled chroma), 'yuv444p' (full-res chroma), 'yuv420p10le' (10-bit per channel), etc. The default (None) selects 'yuv420p' if all shape dimensions are even, else 'yuv444p'. + sandbox_max_run_time_secs: The maximum time in seconds to run the sandbox. + If None, the default limit is 30 minutes. Unused in open source. """ def __init__( @@ -1591,6 +1618,7 @@ def __init__( input_format: str = 'rgb', dtype: _DTypeLike = np.uint8, encoded_format: str | None = None, + sandbox_max_run_time_secs: int | None = None, ) -> None: _check_2d_shape(shape) if fps is None and metadata: @@ -1645,6 +1673,7 @@ def __init__( self.input_format = input_format self.dtype = dtype self.encoded_format = encoded_format + self.sandbox_max_run_time_secs = sandbox_max_run_time_secs if num_rate_specifications == 0 and not ffmpeg_args: qp = 20 if math.prod(self.shape) <= 640 * 480 else 28 self._bitrate_args = ( @@ -1707,6 +1736,7 @@ def __enter__(self) -> 'VideoWriter': stdin=subprocess.PIPE, stderr=subprocess.PIPE, allowed_output_files=[tmp_name], + sandbox_max_run_time_secs=self.sandbox_max_run_time_secs, ) self._proc = self._popen.__enter__() except Exception: diff --git a/mediapy_examples.py b/mediapy_examples.py index 54de193..329b979 100644 --- a/mediapy_examples.py +++ b/mediapy_examples.py @@ -1,4 +1,4 @@ -# Copyright 2025 The mediapy Authors. +# Copyright 2026 The mediapy Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/mediapy_test.py b/mediapy_test.py index 0a54b7a..71a01d2 100755 --- a/mediapy_test.py +++ b/mediapy_test.py @@ -1,4 +1,4 @@ -# Copyright 2025 The mediapy Authors. +# Copyright 2026 The mediapy Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/pdoc_files/make.py b/pdoc_files/make.py index 13ea30a..8a65ac4 100644 --- a/pdoc_files/make.py +++ b/pdoc_files/make.py @@ -1,4 +1,4 @@ -# Copyright 2025 The mediapy Authors. +# Copyright 2026 The mediapy Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License.