From aadcf89253436f0bacbf24070849ce6018cd7f39 Mon Sep 17 00:00:00 2001 From: Craig de Stigter Date: Tue, 2 Jun 2020 14:04:34 +1200 Subject: [PATCH] Add file-like pager: `click.get_pager_file()` --- CHANGES.rst | 2 + docs/api.rst | 2 + docs/utils.rst | 12 ++++ src/click/__init__.py | 1 + src/click/_termui_impl.py | 137 +++++++++++++++++++++----------------- src/click/termui.py | 25 +++++-- 6 files changed, 113 insertions(+), 66 deletions(-) diff --git a/CHANGES.rst b/CHANGES.rst index 4a8e2fe66..a9dc1cd08 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -12,6 +12,8 @@ Unreleased parameter. :issue:`1264`, :pr:`1329` - Add an optional parameter to ``ProgressBar.update`` to set the ``current_item``. :issue:`1226`, :pr:`1332` +- Add ``click.get_pager_file`` for file-like access to an output + pager. :pr:`XXXX` Version 7.1.2 diff --git a/docs/api.rst b/docs/api.rst index 22dd39f4c..0ad5643b7 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -38,6 +38,8 @@ Utilities .. autofunction:: echo_via_pager +.. autofunction:: get_pager_file + .. autofunction:: prompt .. autofunction:: confirm diff --git a/docs/utils.rst b/docs/utils.rst index 6338df941..88780ff0e 100644 --- a/docs/utils.rst +++ b/docs/utils.rst @@ -119,6 +119,18 @@ If you want to use the pager for a lot of text, especially if generating everyth click.echo_via_pager(_generate_output()) +For more complex programs, which can't easily use a simple generator, you +can get access to a writable file-like object for the pager, and write to +that instead: + +.. click:example:: + @click.command() + def less(): + with click.get_pager_file() as pager: + for idx in range(50000): + print(idx, file=pager) + + Screen Clearing --------------- diff --git a/src/click/__init__.py b/src/click/__init__.py index 9cd0129bf..8f7de66ea 100644 --- a/src/click/__init__.py +++ b/src/click/__init__.py @@ -41,6 +41,7 @@ from .termui import confirm from .termui import echo_via_pager from .termui import edit +from .termui import get_pager_file from .termui import get_terminal_size from .termui import getchar from .termui import launch diff --git a/src/click/_termui_impl.py b/src/click/_termui_impl.py index 78372503d..4d8e71e33 100644 --- a/src/click/_termui_impl.py +++ b/src/click/_termui_impl.py @@ -4,6 +4,7 @@ placed in this module and only imported as needed. """ import contextlib +import io import math import os import sys @@ -13,7 +14,6 @@ from ._compat import CYGWIN from ._compat import get_best_encoding from ._compat import isatty -from ._compat import open_stream from ._compat import strip_ansi from ._compat import term_len from ._compat import WIN @@ -329,62 +329,89 @@ def generator(self): self.render_progress() -def pager(generator, color=None): - """Decide what method to use for paging through text.""" +class StripAnsi(io.TextIOWrapper): + @classmethod + def maybe(cls, stream, *, color, encoding): + if not getattr(stream, "encoding", None): + if color: + stream = io.TextIOWrapper(stream, encoding=encoding) + else: + stream = cls(stream) + stream.color = color + return stream + + def write(self, text): + text = strip_ansi(text) + return super().write(text) + + +@contextlib.contextmanager +def get_pager_file(color=None): + """Context manager. + + Yields a writable file-like object which can be used as an output pager. + + .. versionadded:: 8.0 + + :param color: controls if the pager supports ANSI colors or not. The + default is autodetection. + """ stdout = _default_text_stdout() - if not isatty(sys.stdin) or not isatty(stdout): - return _nullpager(stdout, generator, color) pager_cmd = (os.environ.get("PAGER", None) or "").strip() + env = dict(os.environ) if pager_cmd: + # If we're piping to less we might support colors + # if the right flags are passed... + cmd_detail = pager_cmd.rsplit("/", 1)[-1].split() + if color is None and cmd_detail[0] == "less": + less_flags = f"{os.environ.get('LESS', '')}{' '.join(cmd_detail[1:])}" + if not less_flags: + env["LESS"] = "-R" + color = True + elif "r" in less_flags or "R" in less_flags: + color = True + if not isatty(sys.stdin) or not isatty(stdout): + ctx = contextlib.nullcontext((stdout, None)) + elif pager_cmd: if WIN: - return _tempfilepager(generator, pager_cmd, color) - return _pipepager(generator, pager_cmd, color) - if os.environ.get("TERM") in ("dumb", "emacs"): - return _nullpager(stdout, generator, color) - if WIN or sys.platform.startswith("os2"): - return _tempfilepager(generator, "more <", color) - if hasattr(os, "system") and os.system("(less) 2>/dev/null") == 0: - return _pipepager(generator, "less", color) + ctx = _tempfilepager(pager_cmd) + else: + ctx = _pipepager(pager_cmd, env=env) + elif os.environ.get("TERM") in ("dumb", "emacs"): + ctx = contextlib.nullcontext((stdout, None)) + elif WIN or sys.platform.startswith("os2"): + ctx = _tempfilepager("more <") + elif hasattr(os, "system") and os.system("(less) 2>/dev/null") == 0: + ctx = _pipepager("less", env=env) + else: + import tempfile - import tempfile + fd, filename = tempfile.mkstemp() + os.close(fd) + try: + if hasattr(os, "system") and os.system(f'more "{filename}"') == 0: + ctx = _pipepager("more") + else: + ctx = contextlib.nullcontext((stdout, None)) + finally: + os.unlink(filename) - fd, filename = tempfile.mkstemp() - os.close(fd) - try: - if hasattr(os, "system") and os.system(f'more "{filename}"') == 0: - return _pipepager(generator, "more", color) - return _nullpager(stdout, generator, color) - finally: - os.unlink(filename) + with ctx as (stream, encoding): + with StripAnsi.maybe(stream, color=color, encoding=encoding) as text_stream: + yield text_stream -def _pipepager(generator, cmd, color): +@contextlib.contextmanager +def _pipepager(cmd, env=None): """Page through text by feeding it to another program. Invoking a pager through this might support colors. """ import subprocess - env = dict(os.environ) - - # If we're piping to less we might support colors under the - # condition that - cmd_detail = cmd.rsplit("/", 1)[-1].split() - if color is None and cmd_detail[0] == "less": - less_flags = f"{os.environ.get('LESS', '')}{' '.join(cmd_detail[1:])}" - if not less_flags: - env["LESS"] = "-R" - color = True - elif "r" in less_flags or "R" in less_flags: - color = True - c = subprocess.Popen(cmd, shell=True, stdin=subprocess.PIPE, env=env) - encoding = get_best_encoding(c.stdin) try: - for text in generator: - if not color: - text = strip_ansi(text) - - c.stdin.write(text.encode(encoding, "replace")) + encoding = get_best_encoding(c.stdin) + yield c.stdin, encoding except (OSError, KeyboardInterrupt): pass else: @@ -407,30 +434,16 @@ def _pipepager(generator, cmd, color): break -def _tempfilepager(generator, cmd, color): +@contextlib.contextmanager +def _tempfilepager(cmd): """Page through text by invoking a program on a temporary file.""" import tempfile - filename = tempfile.mktemp() - # TODO: This never terminates if the passed generator never terminates. - text = "".join(generator) - if not color: - text = strip_ansi(text) encoding = get_best_encoding(sys.stdout) - with open_stream(filename, "wb")[0] as f: - f.write(text.encode(encoding)) - try: - os.system(f'{cmd} "{filename}"') - finally: - os.unlink(filename) - - -def _nullpager(stream, generator, color): - """Simply print unformatted text. This is the ultimate fallback.""" - for text in generator: - if not color: - text = strip_ansi(text) - stream.write(text) + with tempfile.NamedTemporaryFile(mode="wb") as f: + yield f, encoding + f.flush() + os.system(f'{cmd} "{f.name}"') class Editor: diff --git a/src/click/termui.py b/src/click/termui.py index a1bdf2ab8..cb0703fd6 100644 --- a/src/click/termui.py +++ b/src/click/termui.py @@ -255,6 +255,23 @@ def ioctl_gwinsz(fd): return int(cr[1]), int(cr[0]) +def get_pager_file(color=None): + """Context manager. + + Yields a writable file-like object which can be used as an output pager. + + .. versionadded:: 8.0 + + :param color: controls if the pager supports ANSI colors or not. The + default is autodetection. + """ + from ._termui_impl import get_pager_file + + color = resolve_color_default(color) + + return get_pager_file(color=color) + + def echo_via_pager(text_or_generator, color=None): """This function takes a text and shows it via an environment specific pager on stdout. @@ -267,7 +284,6 @@ def echo_via_pager(text_or_generator, color=None): :param color: controls if the pager supports ANSI colors or not. The default is autodetection. """ - color = resolve_color_default(color) if inspect.isgeneratorfunction(text_or_generator): i = text_or_generator() @@ -279,9 +295,10 @@ def echo_via_pager(text_or_generator, color=None): # convert every element of i to a text type if necessary text_generator = (el if isinstance(el, str) else str(el) for el in i) - from ._termui_impl import pager - - return pager(itertools.chain(text_generator, "\n"), color) + with get_pager_file(color=color) as pager: + for text in itertools.chain(text_generator, "\n"): + # pager.write(text.encode('utf-8')) + pager.write(text) def progressbar(