From e3322e816eaf1ee7e9357c7655bda99bf6e2a9de Mon Sep 17 00:00:00 2001 From: Felix Fontein Date: Sun, 28 Jun 2026 22:03:03 +0200 Subject: [PATCH] Add a new method Session.run_capture() which allows to capture stdout and stderr separately, even while showing the result. --- nox/_cli.py | 3 + nox/command.py | 104 ++++++++++++++++++++- nox/popen.py | 242 ++++++++++++++++++++++++++++++++++++++++++++++-- nox/sessions.py | 194 +++++++++++++++++++++++++++++++++++--- 4 files changed, 520 insertions(+), 23 deletions(-) diff --git a/nox/_cli.py b/nox/_cli.py index ecbabf06..4f583e26 100644 --- a/nox/_cli.py +++ b/nox/_cli.py @@ -17,6 +17,7 @@ from __future__ import annotations __lazy_modules__ = { + "asyncio", "importlib", "importlib.metadata", "nox._options", @@ -36,6 +37,7 @@ "urllib.parse", } +import asyncio import importlib.metadata import os import shutil @@ -316,6 +318,7 @@ def _main(*, main_ep: bool) -> None: ) nox.registry.reset() + asyncio.set_event_loop(asyncio.new_event_loop()) exit_code = execute_workflow(args) # Done; exit. diff --git a/nox/command.py b/nox/command.py index 66af1b52..288ffdcc 100644 --- a/nox/command.py +++ b/nox/command.py @@ -24,13 +24,18 @@ from typing import TYPE_CHECKING, Literal, overload from nox.logger import logger -from nox.popen import DEFAULT_INTERRUPT_TIMEOUT, DEFAULT_TERMINATE_TIMEOUT, popen +from nox.popen import ( + DEFAULT_INTERRUPT_TIMEOUT, + DEFAULT_TERMINATE_TIMEOUT, + popen, + tee_popen, +) if TYPE_CHECKING: from collections.abc import Iterable, Mapping, Sequence from typing import IO -__all__ = ["CommandFailed", "ExternalType", "run", "which"] +__all__ = ["CommandFailed", "ExternalType", "run", "run_capture", "which"] _PLATFORM = sys.platform @@ -216,3 +221,98 @@ def run( raise return output if silent else True + + +def run_capture( + args: Sequence[str | os.PathLike[str]], + *, + env: Mapping[str, str | None] | None = None, + silent: bool = False, + paths: Sequence[str | os.PathLike[str]] | None = None, + success_codes: Iterable[int] | None = None, + success_all: bool = False, + log: bool = True, + external: ExternalType = False, + interrupt_timeout: float | None = DEFAULT_INTERRUPT_TIMEOUT, + terminate_timeout: float | None = DEFAULT_TERMINATE_TIMEOUT, +) -> tuple[str, str, int]: + """Run a command-line program.""" + + if success_all: + if success_codes is not None: + msg = "Cannot use success_all=True with success_codes." + raise ValueError(msg) + success_codes = None + elif success_codes is None: + success_codes = [0] + + cmd, args = args[0], args[1:] + full_cmd = f"{cmd} {_shlex_join(args)}" + + cmd_path = which(os.fspath(cmd), paths) + str_args = [os.fspath(arg) for arg in args] + + if log: + logger.info(full_cmd) + + is_external_tool = paths is not None and not any( + cmd_path.startswith(str(path)) for path in paths + ) + if is_external_tool: + if external == "error": + logger.error( + f"Error: {cmd} is not installed into the virtualenv, it is located" + f" at {cmd_path}. Pass external=True into run() to explicitly allow" + " this." + ) + msg = "External program disallowed." + raise CommandFailed(msg) + if external is False and log: + logger.warning( + f"Warning: {cmd} is not installed into the virtualenv, it is" + f" located at {cmd_path}. This might cause issues! Pass" + " external=True into run() to silence this message." + ) + + env = _clean_env(env) + + try: + if silent: + return_code, stdout_res, stderr_res = popen( + [cmd_path, *str_args], + silent=False, + return_stderr=True, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + interrupt_timeout=interrupt_timeout, + terminate_timeout=terminate_timeout, + ) + else: + return_code, stdout_res, stderr_res = tee_popen( + [cmd_path, *str_args], + env=env, + interrupt_timeout=interrupt_timeout, + terminate_timeout=terminate_timeout, + ) + + if success_codes is not None and return_code not in success_codes: + suffix = ":" if (silent and stderr_res) else "" + logger.error( + f"Command {full_cmd} failed with exit code {return_code}{suffix}" + ) + + if silent and stderr_res: + logger.error(stderr_res) + + msg = f"Returned code {return_code}" + raise CommandFailed(msg) + + if joined_out := "\n".join([out for out in (stderr_res, stdout_res) if out]): + logger.output(joined_out) + + except KeyboardInterrupt: + logger.error("Interrupted...") + raise + + return stdout_res, stderr_res, return_code diff --git a/nox/popen.py b/nox/popen.py index e419e93d..f175f646 100644 --- a/nox/popen.py +++ b/nox/popen.py @@ -14,15 +14,18 @@ from __future__ import annotations -__lazy_modules__ = {"contextlib", "locale"} +__lazy_modules__ = {"asyncio", "contextlib", "locale", "signal"} +import asyncio import contextlib import locale +import signal import subprocess import sys -from typing import IO, TYPE_CHECKING +from typing import IO, TYPE_CHECKING, BinaryIO, Literal, overload if TYPE_CHECKING: + from asyncio.subprocess import Process from collections.abc import Mapping, Sequence __all__ = [ @@ -30,6 +33,7 @@ "DEFAULT_TERMINATE_TIMEOUT", "decode_output", "popen", + "tee_popen", ] @@ -77,16 +81,59 @@ def decode_output(output: bytes) -> str: return output.decode(second_encoding) +@overload +def popen( + args: Sequence[str], + *, + env: Mapping[str, str] | None = ..., + silent: bool = ..., + return_stderr: Literal[False] = ..., + stdout: int | IO[str] | None = ..., + stderr: int | IO[str] | None = ..., + interrupt_timeout: float | None = ..., + terminate_timeout: float | None = ..., +) -> tuple[int, str]: ... + + +@overload +def popen( + args: Sequence[str], + *, + env: Mapping[str, str] | None = ..., + silent: bool = ..., + return_stderr: Literal[True], + stdout: int | IO[str] | None = ..., + stderr: int | IO[str] | None = ..., + interrupt_timeout: float | None = ..., + terminate_timeout: float | None = ..., +) -> tuple[int, str, str]: ... + + +@overload +def popen( + args: Sequence[str], + *, + env: Mapping[str, str] | None = ..., + silent: bool = ..., + return_stderr: bool = ..., + stdout: int | IO[str] | None = ..., + stderr: int | IO[str] | None = ..., + interrupt_timeout: float | None = ..., + terminate_timeout: float | None = ..., +) -> tuple[int, str] | tuple[int, str, str]: ... + + def popen( args: Sequence[str], *, env: Mapping[str, str] | None = None, silent: bool = False, + return_stderr: bool = False, stdout: int | IO[str] | None = None, stderr: int | IO[str] | None = subprocess.STDOUT, interrupt_timeout: float | None = DEFAULT_INTERRUPT_TIMEOUT, terminate_timeout: float | None = DEFAULT_TERMINATE_TIMEOUT, -) -> tuple[int, str]: +) -> tuple[int, str] | tuple[int, str, str]: if silent and stdout is not None: msg = ( "Can not specify silent and stdout; passing a custom stdout always silences" @@ -100,14 +147,197 @@ def popen( proc = subprocess.Popen(args, env=env, stdout=stdout, stderr=stderr) try: - out, _err = proc.communicate() + out, err = proc.communicate() sys.stdout.flush() except KeyboardInterrupt: - out, _err = shutdown_process(proc, interrupt_timeout, terminate_timeout) + out, err = shutdown_process(proc, interrupt_timeout, terminate_timeout) if proc.returncode != 0: raise return_code = proc.wait() - return return_code, decode_output(out) if out else "" + ret_out = decode_output(out) if out else "" + if not return_stderr: + return return_code, ret_out + return return_code, ret_out, decode_output(err) if err else "" + + +class _TeeSubprocess: + def __init__( + self, + *, + loop: asyncio.AbstractEventLoop, + interrupt_timeout: float | None, + terminate_timeout: float | None, + ) -> None: + self.loop = loop + self.stdout: list[bytes] = [] + self.stderr: list[bytes] = [] + self.interrupt_timeout = interrupt_timeout + self.terminate_timeout = terminate_timeout + self.collectors: list[asyncio.Task[None] | asyncio.Task[int]] = [] + self.main_task: asyncio.Task[bool] | None = None + + async def _shutdown_process(self, proc: Process) -> None: + """Gracefully shutdown a child process.""" + _, pending = await asyncio.wait(self.collectors, timeout=self.interrupt_timeout) + if not pending: + return + + proc.terminate() + + _, pending = await asyncio.wait(pending, timeout=self.terminate_timeout) + if not pending: + return + + proc.kill() + + await asyncio.wait(pending) + + async def _main_task(self, proc: Process, wait_task: asyncio.Task[int]) -> bool: + """ + Control interruption flow. + Return ``True`` if the process had a non-zero exit code during shutdown. + """ + try: + await proc.wait() + except asyncio.CancelledError: + # SIGINT causes the task to be cancelled. We first need to uncancel it + # (if possible) since we don't plan to just give up. + if self.main_task and hasattr(self.main_task, "uncancel"): + # Task.uncancel has been added in Python 3.11 + self.main_task.uncancel() + + # Gracefully shutdown procss. + await self._shutdown_process(proc) + + # Return whether at this point, the process' return code is != 0 + # so we can raise KeyboardInterrupt. + return wait_task.result() != 0 + else: + return False + + async def _read_stream( + self, + stream: asyncio.StreamReader, + out_stream: BinaryIO, + out_buffer: list[bytes], + ) -> None: + """ + Read a stream chunk by chunk, append it to ``out_buffer``, and write it to ``out_stream``. + """ + while True: + chunk = await stream.readline() + if len(chunk) == 0: + break + out_buffer.append(chunk) + out_stream.write(chunk) + out_stream.flush() + + async def _stream_subprocess( + self, + args: Sequence[str], + *, + env: Mapping[str, str] | None, + ) -> tuple[bool, int]: + """ + Start the process, all tasks, and wait for them to finish. + Return a tuple ``(is_canceled, return_code)``. + """ + proc = await asyncio.create_subprocess_exec( + *args, + env=env, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + + # These assumptions are true since we pass PIPE to stdout/stderr, + # but mypy doesn't know that. + assert proc.stdout is not None + assert proc.stderr is not None + + # Note that we write everything to sys.stdout, similar to how popen() + # operates by default. + self.collectors.append( + self.loop.create_task( + self._read_stream(proc.stdout, sys.stdout.buffer, self.stdout) + ) + ) + self.collectors.append( + self.loop.create_task( + self._read_stream(proc.stderr, sys.stdout.buffer, self.stderr) + ) + ) + + # The wait task gives us the return code. + wait_task: asyncio.Task[int] = self.loop.create_task(proc.wait()) + self.collectors.append(wait_task) + + # The main task controls the flow (when interrupting). + self.main_task = self.loop.create_task(self._main_task(proc, wait_task)) + + # We wait for all tasks to finish and extract the return code. + await asyncio.wait([*self.collectors, self.main_task]) + return self.main_task.result(), wait_task.result() + + def run( + self, + args: Sequence[str], + *, + env: Mapping[str, str] | None = None, + ) -> tuple[int, bytes, bytes]: + """ + Run the command with the given environment. + Return the process' return code, standard output, and standard error. + + Note that this function is **NOT** thread-safe, since we have to + add a SIGINT handler to the event loop and later remove it. + """ + is_terminating = [False] + + def handle_sigint() -> None: + # Don't handle it twice. + if is_terminating[0]: + return + is_terminating[0] = True + + # Cancel the main task. This triggers graceful shutdown. + if self.main_task: + self.main_task.cancel() + + self.loop.add_signal_handler(signal.SIGINT, handle_sigint) + try: + is_canceled, return_code = self.loop.run_until_complete( + self._stream_subprocess(args, env=env) + ) + finally: + self.loop.remove_signal_handler(signal.SIGINT) + sys.stdout.flush() + if is_canceled: + raise KeyboardInterrupt + return return_code, b"".join(self.stdout), b"".join(self.stderr) + + +def tee_popen( + args: Sequence[str], + *, + env: Mapping[str, str] | None = None, + interrupt_timeout: float | None = DEFAULT_INTERRUPT_TIMEOUT, + terminate_timeout: float | None = DEFAULT_TERMINATE_TIMEOUT, +) -> tuple[int, str, str]: + """ + Run the command with the given environment. + Return a tuple ``(return_code, stdout, stderr)``. + Standard output and standard error are also printed to ``sys.stdout``. + """ + loop = asyncio.get_event_loop() + teesub = _TeeSubprocess( + loop=loop, + interrupt_timeout=interrupt_timeout, + terminate_timeout=terminate_timeout, + ) + return_code, out, err = teesub.run(args, env=env) + ret_out = decode_output(out) + ret_err = decode_output(err) + return return_code, ret_out, ret_err diff --git a/nox/sessions.py b/nox/sessions.py index da534f47..15c6716f 100644 --- a/nox/sessions.py +++ b/nox/sessions.py @@ -551,6 +551,154 @@ def run( terminate_timeout=terminate_timeout, ) + def run_capture( + self, + *args: str | os.PathLike[str], + env: Mapping[str, str | None] | None = None, + include_outer_env: bool = True, + silent: bool = False, + success_codes: Iterable[int] | None = None, + success_all: bool = False, + log: bool = True, + external: ExternalType | None = None, + interrupt_timeout: float | None = DEFAULT_INTERRUPT_TIMEOUT, + terminate_timeout: float | None = DEFAULT_TERMINATE_TIMEOUT, + ) -> tuple[str, str, int] | None: + """Run a command and capture its output and exit code. + + Commands must be specified as a list of strings, for example:: + + stdout, stderr, rc = session.run_capture( + 'pytest', '-k', 'fast', 'tests/') + stdout, stderr, rc = session.run_capture( + 'flake8', '--import-order-style=google') + + The standard output (``stdout``), standard error (``stderr``), + and exit code (``rc``) are always returned. + + You **can not** just pass everything as one string. For example, this + **will not work**:: + + stdout, stderr, rc = session.run_capture('pytest -k fast tests/') + + You can set environment variables for the command using ``env``:: + + stdout, stderr, rc = session.run_capture( + 'bash', '-c', 'echo $SOME_ENV', + env={'SOME_ENV': 'Hello'}) + + You can extend the shutdown timeout to allow long-running cleanup tasks to + complete before being terminated. For example, if you wanted to allow ``pytest`` + extra time to clean up large projects in the case that Nox receives an + interrupt signal from your build system and needs to terminate its child + processes:: + + stdout, stderr, rc = session.run_capture( + 'pytest', '-k', 'long_cleanup', + interrupt_timeout=10.0, + terminate_timeout=2.0) + + You can also tell Nox to treat non-zero exit codes as success using + ``success_codes``. For example, if you wanted to treat the ``pytest`` + "tests discovered, but none selected" error as success:: + + stdout, stderr, rc = session.run_capture( + 'pytest', '-k', 'not slow', + success_codes=[0, 5]) + + If you set ``success_all`` to ``True``, all exit codes will be accepted. + You can + + On Windows, builtin commands like ``del`` cannot be directly invoked, + but you can use ``cmd /c`` to invoke them:: + + stdout, stderr, rc = session.run_capture( + 'cmd', '/c', 'del', 'docs/modules.rst') + + If ``session.run_capture`` fails, it will stop the session and will not run the next steps. + Basically, this will raise a Python exception. Taking this in count, you can use a + ``try...finally`` block for cleanup runs, that will run even if the other runs fail:: + + try: + stdout, stderr, rc = session.run_capture("coverage", "run", "-m", "pytest") + finally: + # Display coverage report even when tests fail. + session.run("coverage", "report") + + If you pass ``silent=True``, the output will only be returned, but not shown to the user. + For example to get the current Git commit ID:: + + stdout, _, _ = session.run_capture( + "git", "rev-parse", "--short", "HEAD", + external=True, silent=True + ) + + print("Current Git commit is", stdout.strip()) + + :param env: A dictionary of environment variables to expose to the + command. By default, all environment variables are passed. You + can block an environment variable from the outer environment by + setting it to None. + :type env: dict or None + :param include_outer_env: Boolean parameter that determines if the + environment variables from the nox invocation environment should + be passed to the command. ``True`` by default. + :type include_outer_env: bool + :param bool silent: Silence command output. In case the command fails, + standard error output will always be shown. ``False`` by default. + :type silent: bool + :param success_codes: A list of return codes that are considered + successful. By default, only ``0`` is considered success. + :type success_codes: list, tuple, or None + :param success_all: If set to ``True``, accept all return codes. + ``False`` by default. + :type success_all: bool + :param external: If False (the default) then programs not in the + virtualenv path will cause a warning. If True, no warning will be + emitted. These warnings can be turned into errors using + ``--error-on-external-run``. This has no effect for sessions that + do not have a virtualenv. + :type external: bool or ``"error"`` + :param interrupt_timeout: The timeout (in seconds) that Nox should wait after it + and its children receive an interrupt signal before sending a terminate + signal to its children. Set to ``None`` to never send a terminate signal. + Default: ``0.3`` + :type interrupt_timeout: float or None + :param terminate_timeout: The timeout (in seconds) that Nox should wait after it + sends a terminate signal to its children before sending a kill signal to + them. Set to ``None`` to never send a kill signal. + Default: ``0.2`` + :type terminate_timeout: float or None + """ + if not args: + msg = "At least one argument required to run_capture()." + raise ValueError(msg) + + if len(args) == 1 and isinstance(args[0], (list, tuple)): + msg = "First argument to `session.run_capture` is a list. Did you mean to use `session.run_capture(*args)`?" + raise ValueError(msg) + + if self._runner.global_config.install_only: + logger.info(f"Skipping {args[0]} run, as --install-only is set.") + return None + + args, env, external = self._prepare_run( + args=args, env=env, include_outer_env=include_outer_env, external=external + ) + + return nox.command.run_capture( + args, + env=env, + paths=self.bin_paths, + silent=silent, + success_codes=success_codes, + success_all=success_all, + log=log, + external=external, + interrupt_timeout=interrupt_timeout, + terminate_timeout=terminate_timeout, + ) + def run_install( self, *args: str | os.PathLike[str], @@ -667,25 +815,16 @@ def run_always( terminate_timeout=terminate_timeout, ) - def _run( + def _prepare_run( self, - *args: str | os.PathLike[str], + *, + args: tuple[str | os.PathLike[str], ...], env: Mapping[str, str | None] | None = None, include_outer_env: bool, - silent: bool, - success_codes: Iterable[int] | None, - log: bool, external: ExternalType | None, - stdout: int | IO[str] | None, - stderr: int | IO[str] | None, - interrupt_timeout: float | None, - terminate_timeout: float | None, - ) -> str | bool: - """Like run(), except that it runs even if --install-only is provided.""" - # Legacy support - run a function given. - if callable(args[0]): - return self._run_func(args[0], args[1:]) # type: ignore[unreachable] - + ) -> tuple[ + tuple[str | os.PathLike[str], ...], Mapping[str, str | None], ExternalType + ]: # Using `"uv"` or `"uvx" when `uv` is the backend is guaranteed to # work, even if it was co-installed with nox. if self.virtualenv.venv_backend == "uv" and nox.virtualenv.UV != "uv": @@ -715,6 +854,31 @@ def _run( if external is None: external = False + return args, env, external + + def _run( + self, + *args: str | os.PathLike[str], + env: Mapping[str, str | None] | None = None, + include_outer_env: bool, + silent: bool, + success_codes: Iterable[int] | None, + log: bool, + external: ExternalType | None, + stdout: int | IO[str] | None, + stderr: int | IO[str] | None, + interrupt_timeout: float | None, + terminate_timeout: float | None, + ) -> str | bool: + """Like run(), except that it runs even if --install-only is provided.""" + # Legacy support - run a function given. + if callable(args[0]): + return self._run_func(args[0], args[1:]) # type: ignore[unreachable] + + args, env, external = self._prepare_run( + args=args, env=env, include_outer_env=include_outer_env, external=external + ) + # Run a shell command. return nox.command.run( args,