Skip to content

Commit

Permalink
Improve echo_via_pager behaviour in face of errors and use of click.e…
Browse files Browse the repository at this point in the history
…cho in the generator

This commit fixes #2674 (echo_via_pager with generators leaves terminal in broken state).

When running the program to reproduce the issue (thanks to @0xDEC0DE for providing it in the issue mentioned above!)

```
import click

def output_generator():
    counter = 0
    while True:
        yield "this is a line of output\n"
        if counter == 1024:
            click.echo("kaboom", err=True)
            click.get_current_context().exit(0)
        counter += 1

@click.command
def kaboom():
    click.echo_via_pager(output_generator)

if __name__ == "__main__":
    kaboom()
```

The "kaboom" message will now be visible immediately (the error will cause click to terminate the pager). ZSH and bash will not be in a broken state (tested using bash and zsh on macos).

Further changes:
- Error handling: terminate the pager and close the file descriptors in a finally block
- KeyboardInterrupt: the pager code will not ignore KeyboardInterrupt anymore. This allows the user to
  search for future output of the generator when using less and then aborting the program using ctrl-c.
- tests: improved the echo_via_pager tests by using a named tuple and separate assertions for pager
  output, stderr and stdout. Added additional test cases.
  • Loading branch information
stefreak committed Sep 7, 2024
1 parent 4a758f5 commit 8be3685
Show file tree
Hide file tree
Showing 3 changed files with 198 additions and 32 deletions.
10 changes: 10 additions & 0 deletions CHANGES.rst
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,16 @@ Unreleased
- When generating a command's name from a decorated function's name, the
suffixes ``_command``, ``_cmd``, ``_group``, and ``_grp`` are removed.
:issue:`2322`
- Improve ``echo_via_pager`` behaviour in face of errors.
:issue:`2674`

- Terminate the pager in case a generator passed to ``echo_via_pager``
raises an exception.
- Ensure to always close the pipe to the pager process and wait for it
to terminate.
- ``echo_via_pager`` will not ignore ``KeyboardInterrupt`` anymore. This
allows the user to search for future output of the generator when
using less and then aborting the program using ctrl-c.


Version 8.1.8
Expand Down
48 changes: 31 additions & 17 deletions src/click/_termui_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -424,26 +424,40 @@ def _pipepager(generator: cabc.Iterable[str], cmd: str, color: bool | None) -> N
text = strip_ansi(text)

stdin.write(text.encode(encoding, "replace"))
except (OSError, KeyboardInterrupt):
except BrokenPipeError:
# In case the pager exited unexpectedly, ignore the broken pipe error.
pass
else:
stdin.close()

# Less doesn't respect ^C, but catches it for its own UI purposes (aborting
# search or other commands inside less).
#
# That means when the user hits ^C, the parent process (click) terminates,
# but less is still alive, paging the output and messing up the terminal.
#
# If the user wants to make the pager exit on ^C, they should set
# `LESS='-K'`. It's not our decision to make.
while True:
except Exception as e:
# In case there is an exception we want to close the pager immediately
# and let the caller handle it.
# Otherwise the pager will keep running, and the user may not notice
# the error message, or worse yet it may leave the terminal in a broken state.
c.terminate()
raise e
finally:
# We must close stdin and wait for the pager to exit before we continue
try:
c.wait()
except KeyboardInterrupt:
stdin.close()
# Close implies flush, so it might throw a BrokenPipeError if the pager
# process exited already.
except BrokenPipeError:
pass
else:
break

# Less doesn't respect ^C, but catches it for its own UI purposes (aborting
# search or other commands inside less).
#
# That means when the user hits ^C, the parent process (click) terminates,
# but less is still alive, paging the output and messing up the terminal.
#
# If the user wants to make the pager exit on ^C, they should set
# `LESS='-K'`. It's not our decision to make.
while True:
try:
c.wait()
except KeyboardInterrupt:
pass
else:
break


def _tempfilepager(generator: cabc.Iterable[str], cmd: str, color: bool | None) -> None:
Expand Down
172 changes: 157 additions & 15 deletions tests/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,9 @@
import pathlib
import stat
import sys
from collections import namedtuple
from io import StringIO
from tempfile import tempdir

import pytest

Expand Down Expand Up @@ -181,32 +183,172 @@ def _test_gen_func():
yield "abc"


def _test_gen_func_fails():
yield "test"
raise RuntimeError("This is a test.")


def _test_gen_func_echo(file=None):
yield "test"
click.echo("hello", file=file)
yield "test"


def _test_simulate_keyboard_interrupt(file=None):
yield "output_before_keyboard_interrupt"
raise KeyboardInterrupt()


EchoViaPagerTest = namedtuple(
"EchoViaPagerTest",
(
"description",
"test_input",
"expected_pager",
"expected_stdout",
"expected_stderr",
"expected_error",
),
)


@pytest.mark.skipif(WIN, reason="Different behavior on windows.")
@pytest.mark.parametrize("cat", ["cat", "cat ", "cat "])
@pytest.mark.parametrize(
"test",
[
# We need lambda here, because pytest will
# reuse the parameters, and then the generators
# are already used and will not yield anymore
("just text\n", lambda: "just text"),
("iterable\n", lambda: ["itera", "ble"]),
("abcabc\n", lambda: _test_gen_func),
("abcabc\n", lambda: _test_gen_func()),
("012345\n", lambda: (c for c in range(6))),
# We need to pass a parameter function instead of a plain param
# as pytest.mark.parametrize will reuse the parameters causing the
# generators to be used up so they will not yield anymore
EchoViaPagerTest(
description="Plain string argument",
test_input=lambda: "just text",
expected_pager="just text\n",
expected_stdout="",
expected_stderr="",
expected_error=None,
),
EchoViaPagerTest(
description="Iterable argument",
test_input=lambda: ["itera", "ble"],
expected_pager="iterable\n",
expected_stdout="",
expected_stderr="",
expected_error=None,
),
EchoViaPagerTest(
description="Generator function argument",
test_input=lambda: _test_gen_func,
expected_pager="abcabc\n",
expected_stdout="",
expected_stderr="",
expected_error=None,
),
EchoViaPagerTest(
description="String generator argument",
test_input=lambda: _test_gen_func(),
expected_pager="abcabc\n",
expected_stdout="",
expected_stderr="",
expected_error=None,
),
EchoViaPagerTest(
description="Number generator expression argument",
test_input=lambda: (c for c in range(6)),
expected_pager="012345\n",
expected_stdout="",
expected_stderr="",
expected_error=None,
),
EchoViaPagerTest(
description="Exception in generator function argument",
test_input=lambda: _test_gen_func_fails,
# Because generator throws early on, the pager did not have
# a chance yet to write the file.
expected_pager=None,
expected_stdout="",
expected_stderr="",
expected_error=RuntimeError,
),
EchoViaPagerTest(
description="Exception in generator argument",
test_input=lambda: _test_gen_func_fails,
# Because generator throws early on, the pager did not have a
# chance yet to write the file.
expected_pager=None,
expected_stdout="",
expected_stderr="",
expected_error=RuntimeError,
),
EchoViaPagerTest(
description="Keyboard interrupt should not terminate the pager",
test_input=lambda: _test_simulate_keyboard_interrupt(),
# Due to the keyboard interrupt during pager execution, click program
# should abort, but the pager should stay open.
# This allows users to cancel the program and search in the pager
# output, before they decide to terminate the pager.
expected_pager="output_before_keyboard_interrupt",
expected_stdout="",
expected_stderr="",
expected_error=KeyboardInterrupt,
),
EchoViaPagerTest(
description="Writing to stdout during generator execution",
test_input=lambda: _test_gen_func_echo(),
expected_pager="testtest\n",
expected_stdout="hello\n",
expected_stderr="",
expected_error=None,
),
EchoViaPagerTest(
description="Writing to stderr during generator execution",
test_input=lambda: _test_gen_func_echo(file=sys.stderr),
expected_pager="testtest\n",
expected_stdout="",
expected_stderr="hello\n",
expected_error=None,
),
],
)
def test_echo_via_pager(monkeypatch, capfd, cat, test):
monkeypatch.setitem(os.environ, "PAGER", cat)
def test_echo_via_pager(monkeypatch, capfd, test):
pager_out_tmp = f"{tempdir}/pager_out.txt"

if os.path.exists(pager_out_tmp):
os.remove(pager_out_tmp)

monkeypatch.setitem(os.environ, "PAGER", f"cat > {pager_out_tmp}")
monkeypatch.setattr(click._termui_impl, "isatty", lambda x: True)

expected_output = test[0]
test_input = test[1]()
test_input = test.test_input()
expected_pager = test.expected_pager
expected_stdout = test.expected_stdout
expected_stderr = test.expected_stderr
expected_error = test.expected_error

click.echo_via_pager(test_input)
if expected_error:
with pytest.raises(expected_error):
click.echo_via_pager(test_input)
else:
click.echo_via_pager(test_input)

out, err = capfd.readouterr()
assert out == expected_output

if os.path.exists(pager_out_tmp):
with open(pager_out_tmp) as f:
pager = f.read()
else:
# The pager process was not started or has been
# terminated before it could finish writing
pager = None

assert (
pager == expected_pager
), f"Unexpected pager output in test case '{test.description}'"
assert (
out == expected_stdout
), f"Unexpected stdout in test case '{test.description}'"
assert (
err == expected_stderr
), f"Unexpected stderr in test case '{test.description}'"


@pytest.mark.skipif(WIN, reason="Test does not make sense on Windows.")
Expand Down

0 comments on commit 8be3685

Please sign in to comment.