diff --git a/src/ffmpeg_normalize/__init__.py b/src/ffmpeg_normalize/__init__.py index 96a9ad0..2815ef3 100644 --- a/src/ffmpeg_normalize/__init__.py +++ b/src/ffmpeg_normalize/__init__.py @@ -1,5 +1,6 @@ import importlib.metadata +from ._cmd_utils import ffmpeg_env from ._errors import FFmpegNormalizeError from ._ffmpeg_normalize import FFmpegNormalize from ._media_file import MediaFile @@ -17,5 +18,6 @@ "VideoStream", "SubtitleStream", "MediaStream", + "ffmpeg_env", "__version__", ] diff --git a/src/ffmpeg_normalize/_cmd_utils.py b/src/ffmpeg_normalize/_cmd_utils.py index 352496c..e491c40 100644 --- a/src/ffmpeg_normalize/_cmd_utils.py +++ b/src/ffmpeg_normalize/_cmd_utils.py @@ -1,10 +1,12 @@ from __future__ import annotations +import contextvars import logging import os import re import shlex import subprocess +from contextlib import contextmanager from shutil import which from typing import Any, Iterator @@ -14,6 +16,30 @@ _logger = logging.getLogger(__name__) +_ffmpeg_env_var: contextvars.ContextVar[dict[str, str] | None] = contextvars.ContextVar( + "ffmpeg_env", default=None +) + + +@contextmanager +def ffmpeg_env(env: dict[str, str] | None) -> Iterator[None]: + """ + Temporarily set the environment for subprocess.Popen. + + Args: + env: Environment dict to pass to subprocess.Popen. + """ + token = _ffmpeg_env_var.set(env) + try: + yield + finally: + _ffmpeg_env_var.reset(token) + + +def _get_ffmpeg_env() -> dict[str, str] | None: + return _ffmpeg_env_var.get() + + DUR_REGEX = re.compile( r"Duration: (?P\d{2}):(?P\d{2}):(?P\d{2})\.(?P\d{2})" ) @@ -76,7 +102,9 @@ def run_ffmpeg_command(self, cmd: list[str]) -> Iterator[float]: # wrapper for 'ffmpeg-progress-yield' _logger.debug(f"Running command: {shlex.join(cmd)}") with FfmpegProgress(cmd, dry_run=self.dry) as ff: - yield from ff.run_command_with_progress() + yield from ff.run_command_with_progress( + popen_kwargs={"env": _get_ffmpeg_env()} + ) self.output = ff.stderr @@ -107,6 +135,7 @@ def run_command(self, cmd: list[str]) -> CommandRunner: stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=False, + env=_get_ffmpeg_env(), ) stdout_bytes, stderr_bytes = p.communicate() diff --git a/tests/test_all.py b/tests/test_all.py index b0b78e0..0feaa61 100644 --- a/tests/test_all.py +++ b/tests/test_all.py @@ -658,3 +658,27 @@ def test_valid_file_passes_validation(self): # Should succeed (file exists and has audio) ffmpeg_normalize_call(["tests/test.mp4", "-n"]) # dry run # No assertion needed - if validation fails, the command will error + + +def test_ffmpeg_env(): + """Verify that ffmpeg_env context manager sets environment correctly.""" + from ffmpeg_normalize._cmd_utils import ffmpeg_env, _get_ffmpeg_env, CommandRunner + + original_env = _get_ffmpeg_env() + assert original_env is None + + test_env = os.environ.copy() + test_env.update({"TEST_FFMPEG_ENV_VAR": "12345"}) + + with ffmpeg_env(test_env): + env_in_context = _get_ffmpeg_env() + assert env_in_context == test_env + # Check that commandrunner uses the env + cmd = CommandRunner().run_command( + [sys.executable, "-c", "import os; print(os.getenv('TEST_FFMPEG_ENV_VAR'))"] + ) + assert cmd.get_output().strip() == "12345" + + # After context, environment should be reset + env_after = _get_ffmpeg_env() + assert env_after is None