Refactor output printer so that it can render big outputs without mem…

…ory issues.

Previously, an expression like `b'\x90' * 40_000_000` would kill ptpython
because it rendered the whole output at once. This implementation streams the
rendering logic while it's paginating.

ptpython/printer.py

@@ -0,0 +1,374 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Generator, Iterable
+
+from prompt_toolkit.formatted_text import (
+    HTML,
+    AnyFormattedText,
+    FormattedText,
+    OneStyleAndTextTuple,
+    StyleAndTextTuples,
+    fragment_list_width,
+    merge_formatted_text,
+    to_formatted_text,
+)
+from prompt_toolkit.formatted_text.utils import split_lines
+from prompt_toolkit.key_binding import KeyBindings, KeyPressEvent
+from prompt_toolkit.output import Output
+from prompt_toolkit.shortcuts import PromptSession, print_formatted_text
+from prompt_toolkit.styles import BaseStyle, StyleTransformation
+from prompt_toolkit.styles.pygments import pygments_token_to_classname
+from prompt_toolkit.utils import get_cwidth
+from pygments.lexers import PythonLexer
+
+__all__ = ["OutputPrinter"]
+
+# Never reformat results larger than this:
+MAX_REFORMAT_SIZE = 1_000_000
+
+
+@dataclass
+class OutputPrinter:
+    """
+    Result printer.
+
+    Usage::
+
+        printer = OutputPrinter(...)
+        printer.display_result(...)
+        printer.display_exception(...)
+    """
+
+    output: Output
+    style: BaseStyle
+    title: AnyFormattedText
+    style_transformation: StyleTransformation
+
+    def display_result(
+        self,
+        result: object,
+        *,
+        out_prompt: AnyFormattedText,
+        reformat: bool,
+        highlight: bool,
+        paginate: bool,
+    ) -> None:
+        """
+        Show __repr__ for an `eval` result and print to output.
+
+        :param reformat: Reformat result using 'black' before printing if the
+            result is parsable as Python code.
+        :param highlight: Syntax highlight the result.
+        :param paginate: Show paginator when the result does not fit on the
+            screen.
+        """
+        out_prompt = to_formatted_text(out_prompt)
+        out_prompt_width = fragment_list_width(out_prompt)
+
+        text_output = self._insert_out_prompt_and_split_lines(
+            self._format_result_output(
+                result,
+                reformat=reformat,
+                highlight=highlight,
+                line_length=self.output.get_size().columns - out_prompt_width,
+            ),
+            out_prompt,
+        )
+
+        if paginate:
+            text_output = self._apply_soft_wrapping(text_output, out_prompt)
+            self._print_paginated_formatted_text(text_output)
+        else:
+            for line in text_output:
+                self._print_formatted_text(line)
+
+        self.output.flush()
+
+    def display_exception(self, e: BaseException) -> None:
+        tokens = self._format_exception_output(e)
+
+        print_formatted_text(
+            tokens,
+            style=self._current_style,
+            style_transformation=self.style_transformation,
+            include_default_pygments_style=False,
+            output=self.output,
+        )
+        self.output.flush()
+
+    def _print_formatted_text(self, line: StyleAndTextTuples, end: str = "\n") -> None:
+        print_formatted_text(
+            FormattedText(line),
+            style=self.style,
+            style_transformation=self.style_transformation,
+            include_default_pygments_style=False,
+            output=self.output,
+            end=end,
+        )
+
+    def _format_result_output(
+        self,
+        result: object,
+        *,
+        reformat: bool,
+        highlight: bool,
+        line_length: int,
+    ) -> Generator[OneStyleAndTextTuple, None, None]:
+        """
+        Format __repr__ for an `eval` result.
+
+        Note: this can raise `KeyboardInterrupt` if either calling `__repr__`,
+              `__pt_repr__` or formatting the output with "Black" takes to long
+              and the user presses Control-C.
+        """
+        # If __pt_repr__ is present, take this. This can return prompt_toolkit
+        # formatted text.
+        try:
+            if hasattr(result, "__pt_repr__"):
+                formatted_result_repr = to_formatted_text(
+                    getattr(result, "__pt_repr__")()
+                )
+                yield from formatted_result_repr
+                return
+        except KeyboardInterrupt:
+            raise  # Don't catch here.
+        except:
+            # For bad code, `__getattr__` can raise something that's not an
+            # `AttributeError`. This happens already when calling `hasattr()`.
+            pass
+
+        # Call `__repr__` of given object first, to turn it in a string.
+        try:
+            result_repr = repr(result)
+        except KeyboardInterrupt:
+            raise  # Don't catch here.
+        except BaseException as e:
+            # Calling repr failed.
+            self.display_exception(e)
+            return
+
+        # Determine whether it's valid Python code. If not,
+        # reformatting/highlighting won't be applied.
+        if len(result_repr) < MAX_REFORMAT_SIZE:
+            try:
+                compile(result_repr, "", "eval")
+            except SyntaxError:
+                valid_python = False
+            else:
+                valid_python = True
+        else:
+            valid_python = False
+
+        if valid_python and reformat:
+            # Inline import. Slightly speed up start-up time if black is
+            # not used.
+            try:
+                import black
+
+                if not hasattr(black, "Mode"):
+                    raise ImportError
+            except ImportError:
+                pass  # no Black package in your installation
+            else:
+                result_repr = black.format_str(
+                    result_repr,
+                    mode=black.Mode(line_length=line_length),
+                )
+
+        if valid_python and highlight:
+            yield from _lex_python_result(result_repr)
+        else:
+            yield ("", result_repr)
+
+    def _insert_out_prompt_and_split_lines(
+        self, result: Iterable[OneStyleAndTextTuple], out_prompt: AnyFormattedText
+    ) -> Iterable[StyleAndTextTuples]:
+        r"""
+        Split styled result in lines (based on the \n characters in the result)
+        an insert output prompt on whitespace in front of each line. (This does
+        not yet do the soft wrapping.)
+
+        Yield lines as a result.
+        """
+        out_prompt = to_formatted_text(out_prompt)
+        out_prompt_width = fragment_list_width(out_prompt)
+        prefix = ("", " " * out_prompt_width)
+
+        for i, line in enumerate(split_lines(result)):
+            if i == 0:
+                line = [*out_prompt, *line]
+            else:
+                line = [prefix, *line]
+            yield line
+
+    def _apply_soft_wrapping(
+        self, lines: Iterable[StyleAndTextTuples], out_prompt: AnyFormattedText
+    ) -> Iterable[StyleAndTextTuples]:
+        """
+        Apply soft wrapping to the given lines. Wrap according to the terminal
+        width. Insert whitespace in front of each wrapped line to align it with
+        the output prompt.
+        """
+        out_prompt = to_formatted_text(out_prompt)
+        out_prompt_width = fragment_list_width(to_formatted_text(out_prompt))
+        prefix = ("", " " * out_prompt_width)
+        line_length = self.output.get_size().columns - out_prompt_width
+
+        # Iterate over hard wrapped lines.
+        for lineno, line in enumerate(lines):
+            columns_in_buffer = 0
+            current_line: list[OneStyleAndTextTuple] = []
+
+            for style, text, *_ in line:
+                for c in text:
+                    width = get_cwidth(c)
+
+                    # (Soft) wrap line if it doesn't fit.
+                    if columns_in_buffer + width > line_length:
+                        yield current_line
+                        columns_in_buffer = 0
+                        current_line = [prefix]
+
+                    columns_in_buffer += width
+                    current_line.append((style, c))
+
+            if len(current_line) > 0:
+                yield current_line
+
+    def _print_paginated_formatted_text(
+        self, lines: Iterable[StyleAndTextTuples]
+    ) -> None:
+        """
+        Print formatted text, using --MORE-- style pagination.
+        (Avoid filling up the terminal's scrollback buffer.)
+        """
+        pager_prompt = create_pager_prompt(self.style, self.title)
+
+        abort = False
+        print_all = False
+
+        # Max number of lines allowed in the buffer before painting.
+        size = self.output.get_size()
+        max_rows = size.rows - 1
+
+        # Page buffer.
+        page: StyleAndTextTuples = []
+
+        def show_pager() -> None:
+            nonlocal abort, max_rows, print_all
+
+            # Run pager prompt in another thread.
+            # Same as for the input. This prevents issues with nested event
+            # loops.
+            pager_result = pager_prompt.prompt(in_thread=True)
+
+            if pager_result == PagerResult.ABORT:
+                print("...")
+                abort = True
+
+            elif pager_result == PagerResult.NEXT_LINE:
+                max_rows = 1
+
+            elif pager_result == PagerResult.NEXT_PAGE:
+                max_rows = size.rows - 1
+
+            elif pager_result == PagerResult.PRINT_ALL:
+                print_all = True
+
+        # Loop over lines. Show --MORE-- prompt when page is filled.
+        rows = 0
+
+        for lineno, line in enumerate(lines):
+            page.extend(line)
+            page.append(("", "\n"))
+            rows += 1
+
+            if rows >= max_rows:
+                self._print_formatted_text(page, end="")
+                page = []
+                rows = 0
+
+                if not print_all:
+                    show_pager()
+                    if abort:
+                        return
+
+        self._print_formatted_text(page)
+
+
+class PagerResult(Enum):
+    ABORT = "ABORT"
+    NEXT_LINE = "NEXT_LINE"
+    NEXT_PAGE = "NEXT_PAGE"
+    PRINT_ALL = "PRINT_ALL"
+
+
+def create_pager_prompt(
+    style: BaseStyle, title: AnyFormattedText = ""
+) -> PromptSession[PagerResult]:
+    """
+    Create a "--MORE--" prompt for paginated output.
+    """
+    bindings = KeyBindings()
+
+    @bindings.add("enter")
+    @bindings.add("down")
+    def next_line(event: KeyPressEvent) -> None:
+        event.app.exit(result=PagerResult.NEXT_LINE)
+
+    @bindings.add("space")
+    def next_page(event: KeyPressEvent) -> None:
+        event.app.exit(result=PagerResult.NEXT_PAGE)
+
+    @bindings.add("a")
+    def print_all(event: KeyPressEvent) -> None:
+        event.app.exit(result=PagerResult.PRINT_ALL)
+
+    @bindings.add("q")
+    @bindings.add("c-c")
+    @bindings.add("c-d")
+    @bindings.add("escape", eager=True)
+    def no(event: KeyPressEvent) -> None:
+        event.app.exit(result=PagerResult.ABORT)
+
+    @bindings.add("<any>")
+    def _(event: KeyPressEvent) -> None:
+        "Disallow inserting other text."
+        pass
+
+    style
+
+    session: PromptSession[PagerResult] = PromptSession(
+        merge_formatted_text(
+            [
+                title,
+                HTML(
+                    "<status-toolbar>"
+                    "<more> -- MORE -- </more> "
+                    "<key>[Enter]</key> Scroll "
+                    "<key>[Space]</key> Next page "
+                    "<key>[a]</key> Print all "
+                    "<key>[q]</key> Quit "
+                    "</status-toolbar>: "
+                ),
+            ]
+        ),
+        key_bindings=bindings,
+        erase_when_done=True,
+        style=style,
+    )
+    return session
+
+
+def _lex_python_result(result: str) -> Generator[tuple[str, str], None, None]:
+    "Return token list for Python string."
+    lexer = PythonLexer()
+    # Use `get_tokens_unprocessed`, so that we get exactly the same string,
+    # without line endings appended. `print_formatted_text` already appends a
+    # line ending, and otherwise we'll have two line endings.
+    tokens = lexer.get_tokens_unprocessed(result)
+
+    for index, tokentype, text in tokens:
+        yield ("class:" + pygments_token_to_classname(tokentype), text)