From be1b14d42b55ff4040dcca294ec67465a5968fea Mon Sep 17 00:00:00 2001 From: Will McGugan Date: Thu, 22 Oct 2020 21:17:03 +0100 Subject: [PATCH] indentaion guides --- rich/default_styles.py | 1 + rich/pretty.py | 13 +++++- rich/progress.py | 3 +- rich/scope.py | 12 +++++- rich/syntax.py | 32 +++++++++++--- rich/text.py | 84 ++++++++++++++++++++++++++++++++++++- rich/traceback.py | 5 +++ tests/test_pretty.py | 18 +++++++- tests/test_syntax.py | 20 +++++++++ tests/test_text.py | 38 +++++++++++++++++ tools/stress_test_pretty.py | 2 +- 11 files changed, 214 insertions(+), 14 deletions(-) diff --git a/rich/default_styles.py b/rich/default_styles.py index 94c106cc..fccf8694 100644 --- a/rich/default_styles.py +++ b/rich/default_styles.py @@ -55,6 +55,7 @@ DEFAULT_STYLES: Dict[str, Style] = { "log.time": Style(color="cyan", dim=True), "log.message": Style.null(), "log.path": Style(dim=True), + "repr.indent": Style(color="green", dim=True), "repr.error": Style(color="red", bold=True), "repr.str": Style(color="green", italic=False, bold=False), "repr.brace": Style(bold=True), diff --git a/rich/pretty.py b/rich/pretty.py index 2e334397..0bf5a4be 100644 --- a/rich/pretty.py +++ b/rich/pretty.py @@ -40,6 +40,7 @@ def install( console: "Console" = None, overflow: "OverflowMethod" = "ignore", crop: bool = False, + indent_guides: bool = False, ) -> None: """Install automatic pretty printing in the Python REPL. @@ -47,6 +48,7 @@ def install( console (Console, optional): Console instance or ``None`` to use global console. Defaults to None. overflow (Optional[OverflowMethod], optional): Overflow method. Defaults to "ignore". crop (Optional[bool], optional): Enable cropping of long lines. Defaults to False. + indent_guides (bool, optional): Enable indentation guides. Defaults to False. """ from rich import get_console @@ -60,7 +62,7 @@ def install( console.print( value if hasattr(value, "__rich_console__") or hasattr(value, "__rich__") - else Pretty(value, overflow=overflow), + else Pretty(value, overflow=overflow, indent_guides=indent_guides), crop=crop, ) builtins._ = value # type: ignore @@ -78,6 +80,7 @@ class Pretty: justify (JustifyMethod, optional): Justify method, or None for default. Defaults to None. overflow (OverflowMethod, optional): Overflow method, or None for default. Defaults to None. no_wrap (Optional[bool], optional): Disable word wrapping. Defaults to False. + indent_guides (bool, optional): Enable indentation guides. Defaults to False. """ def __init__( @@ -89,6 +92,7 @@ class Pretty: justify: "JustifyMethod" = None, overflow: Optional["OverflowMethod"] = "crop", no_wrap: Optional[bool] = False, + indent_guides: bool = False, ) -> None: self._object = _object self.highlighter = highlighter or ReprHighlighter() @@ -96,6 +100,7 @@ class Pretty: self.justify = justify self.overflow = overflow self.no_wrap = no_wrap + self.indent_guides = indent_guides def __rich_console__( self, console: "Console", options: "ConsoleOptions" @@ -111,6 +116,10 @@ class Pretty: style="pretty", ) pretty_text = self.highlighter(pretty_text) + if self.indent_guides and not options.ascii_only: + pretty_text = pretty_text.with_indent_guides( + self.indent_size, style="repr.indent" + ) yield pretty_text def __rich_measure__(self, console: "Console", max_width: int) -> "Measurement": @@ -415,4 +424,4 @@ if __name__ == "__main__": # pragma: no cover from rich import print - print(Pretty(data)) + print(Pretty(data, indent_guides=True)) diff --git a/rich/progress.py b/rich/progress.py index 8db3bb72..7a3ad0cb 100644 --- a/rich/progress.py +++ b/rich/progress.py @@ -780,7 +780,8 @@ class Progress(JupyterMixin, RenderHook): popleft() while len(_progress) > 1000: popleft() - _progress.append(ProgressSample(current_time, update_completed)) + if update_completed > 0: + _progress.append(ProgressSample(current_time, update_completed)) def reset( self, diff --git a/rich/scope.py b/rich/scope.py index 8456eb2a..725ffa27 100644 --- a/rich/scope.py +++ b/rich/scope.py @@ -14,7 +14,11 @@ if TYPE_CHECKING: def render_scope( - scope: Mapping, *, title: TextType = None, sort_keys: bool = True + scope: Mapping, + *, + title: TextType = None, + sort_keys: bool = True, + indent_guides: bool = False ) -> "ConsoleRenderable": """Render python variables in a given scope. @@ -22,6 +26,7 @@ def render_scope( scope (Mapping): A mapping containing variable names and values. title (str, optional): Optional title. Defaults to None. sort_keys (bool, optional): Enable sorting of items. Defaults to True. + indent_guides (bool, optional): Enable indentaton guides. Defaults to False. Returns: RenderableType: A renderable object. @@ -41,7 +46,10 @@ def render_scope( (key, "scope.key.special" if key.startswith("__") else "scope.key"), (" =", "scope.equals"), ) - items_table.add_row(key_text, Pretty(value, highlighter=highlighter)) + items_table.add_row( + key_text, + Pretty(value, highlighter=highlighter, indent_guides=indent_guides), + ) return Panel.fit( items_table, title=title, diff --git a/rich/syntax.py b/rich/syntax.py index 670ddba7..36e295ea 100644 --- a/rich/syntax.py +++ b/rich/syntax.py @@ -203,6 +203,7 @@ class Syntax(JupyterMixin): tab_size (int, optional): Size of tabs. Defaults to 4. word_wrap (bool, optional): Enable word wrapping. background_color (str, optional): Optional background color, or None to use theme color. Defaults to None. + indent_guides (bool, optional): Show indent guides. Defaults to False. """ _pygments_style_class: Type[PygmentsStyle] @@ -235,6 +236,7 @@ class Syntax(JupyterMixin): tab_size: int = 4, word_wrap: bool = False, background_color: str = None, + indent_guides: bool = False, ) -> None: self.code = code self.lexer_name = lexer_name @@ -247,6 +249,7 @@ class Syntax(JupyterMixin): self.tab_size = tab_size self.word_wrap = word_wrap self.background_color = background_color + self.indent_guides = indent_guides self._theme = self.get_theme(theme) @@ -265,6 +268,7 @@ class Syntax(JupyterMixin): tab_size: int = 4, word_wrap: bool = False, background_color: str = None, + indent_guides: bool = False, ) -> "Syntax": """Construct a Syntax object from a file. @@ -281,6 +285,7 @@ class Syntax(JupyterMixin): tab_size (int, optional): Size of tabs. Defaults to 4. word_wrap (bool, optional): Enable word wrapping of code. background_color (str, optional): Optional background color, or None to use theme color. Defaults to None. + indent_guides (bool, optional): Show indent guides. Defaults to False. Returns: [Syntax]: A Syntax object that may be printed to the console @@ -318,6 +323,7 @@ class Syntax(JupyterMixin): tab_size=tab_size, word_wrap=word_wrap, background_color=background_color, + indent_guides=indent_guides, ) def _get_base_style(self) -> Style: @@ -439,6 +445,17 @@ class Syntax(JupyterMixin): text = self.highlight(code) text.remove_suffix("\n") text.expand_tabs(self.tab_size) + + ( + background_style, + number_style, + highlight_number_style, + ) = self._get_number_styles(console) + + if self.indent_guides and not options.ascii_only: + style = self._get_base_style() + self._theme.get_style_for_token(Comment) + text = text.with_indent_guides(self.tab_size, style=style) + if not self.line_numbers: # Simple case of just rendering text yield from console.render(text, options=options.update(width=code_width)) @@ -454,12 +471,6 @@ class Syntax(JupyterMixin): numbers_column_width = self._numbers_column_width render_options = options.update(width=code_width) - ( - background_style, - number_style, - highlight_number_style, - ) = self._get_number_styles(console) - highlight_line = self.highlight_lines.__contains__ _Segment = Segment padding = _Segment(" " * numbers_column_width + " ", background_style) @@ -519,6 +530,14 @@ if __name__ == "__main__": # pragma: no cover default=None, help="force color for non-terminals", ) + parser.add_argument( + "-i", + "--indent-guides", + dest="indent_guides", + action="store_true", + default=False, + help="display indent guides", + ) parser.add_argument( "-l", "--line-numbers", @@ -572,5 +591,6 @@ if __name__ == "__main__": # pragma: no cover word_wrap=args.word_wrap, theme=args.theme, background_color=args.background_color, + indent_guides=args.indent_guides, ) console.print(syntax, soft_wrap=args.soft_wrap) diff --git a/rich/text.py b/rich/text.py index 8e3d5380..d9a5845a 100644 --- a/rich/text.py +++ b/rich/text.py @@ -1,4 +1,5 @@ -from functools import partial +from functools import partial, reduce +from math import gcd import re from operator import itemgetter from typing import ( @@ -1008,6 +1009,75 @@ class Text(JupyterMixin): append(line) return lines + def detect_indentation(self) -> int: + """Auto-detect indentation of code. + + Returns: + int: Number of spaces used to indent code. + """ + + _indentations = { + len(match.group(1)) + for match in re.finditer(r"^( *)(.*)$", self.plain, flags=re.MULTILINE) + } + + try: + indentation = ( + reduce(gcd, [indent for indent in _indentations if not indent % 2]) or 1 + ) + except TypeError: + indentation = 1 + + return indentation + + def with_indent_guides( + self, + indent_size: int = None, + *, + character: str = "│", + style: StyleType = "dim green", + ) -> "Text": + """Adds indent guide lines to text. + + Args: + indent_size (int): Size of indentation. + character (str, optional): Character to use for indentation. Defaults to "│". + style (Union[Style, str], optional): Style of indent guides. + + Returns: + [type]: [description] + """ + + _indent_size = self.detect_indentation() if indent_size is None else indent_size + + text = self.copy() + text.expand_tabs() + indent_line = f"{character}{' ' * (_indent_size - 1)}" + + re_indent = re.compile(r"^( *)(.*)$") + new_lines: List[Text] = [] + add_line = new_lines.append + blank_lines = 0 + for line in text.split(): + match = re_indent.match(line.plain) + if not match.group(2): + blank_lines += 1 + continue + indent = match.group(1) + full_indents, remaining_space = divmod(len(indent), _indent_size) + new_indent = f"{indent_line * full_indents}{' ' * remaining_space}" + line.plain = new_indent + line.plain[len(new_indent) :] + line.stylize(style, 0, len(new_indent)) + if blank_lines: + new_lines.extend([Text(new_indent, style=style)] * blank_lines) + blank_lines = 0 + add_line(line) + if blank_lines: + new_lines.extend([Text("", style=style)] * blank_lines) + + new_text = Text("\n").join(new_lines) + return new_text + if __name__ == "__main__": # pragma: no cover from rich import print @@ -1015,3 +1085,15 @@ if __name__ == "__main__": # pragma: no cover text = Text("\n\tHello\n") text.expand_tabs(4) print(text) + + code = """ +def __add__(self, other: Any) -> "Text": + if isinstance(other, (str, Text)): + result = self.copy() + result.append(other) + return result + return NotImplemented +""" + text = Text(code) + text = text.with_indent_guides() + print(text) \ No newline at end of file diff --git a/rich/traceback.py b/rich/traceback.py index 71ce631b..62c28fc0 100644 --- a/rich/traceback.py +++ b/rich/traceback.py @@ -3,12 +3,14 @@ from __future__ import absolute_import import platform import sys from dataclasses import dataclass, field +from textwrap import indent from traceback import walk_tb from types import TracebackType from typing import Callable, Dict, List, Optional, Type from pygments.lexers import guess_lexer_for_filename from pygments.token import ( + Comment, Keyword, Name, Number, @@ -279,6 +281,7 @@ class Traceback: "pygments.string": token_style(String), "pygments.function": token_style(Name.Function), "pygments.number": token_style(Number), + "repr.indent": token_style(Comment), "repr.str": token_style(String), "repr.brace": token_style(TextToken) + Style(bold=True), "repr.number": token_style(Number), @@ -407,6 +410,7 @@ class Traceback: highlight_lines={frame.lineno}, word_wrap=self.word_wrap, code_width=88, + indent_guides=False, ) yield "" except Exception: @@ -437,6 +441,7 @@ if __name__ == "__main__": # pragma: no cover print(one / a) def foo(a): + zed = { "characters": { "Paul Atriedies", diff --git a/tests/test_pretty.py b/tests/test_pretty.py index fdf7e749..bfce2312 100644 --- a/tests/test_pretty.py +++ b/tests/test_pretty.py @@ -4,7 +4,7 @@ import io import sys from rich.console import Console -from rich.pretty import install, pretty_repr, Node +from rich.pretty import install, Pretty, pretty_repr, Node def test_install(): @@ -77,3 +77,19 @@ def test_tuple_of_one(): def test_node(): node = Node("abc") assert pretty_repr(node) == "abc: " + + +def test_indent_lines(): + console = Console(width=100, color_system=None) + console.begin_capture() + console.print(Pretty([100, 200], indent_guides=True), width=8) + expected = """\ +[ +│ 100, +│ 200 +] +""" + result = console.end_capture() + print(repr(result)) + print(result) + assert result == expected diff --git a/tests/test_syntax.py b/tests/test_syntax.py index c63b0f23..4ffad6f1 100644 --- a/tests/test_syntax.py +++ b/tests/test_syntax.py @@ -54,6 +54,26 @@ def test_python_render(): assert rendered_syntax == expected +def test_python_render_indent_guides(): + syntax = Panel.fit( + Syntax( + CODE, + lexer_name="python", + line_numbers=True, + line_range=(2, 10), + theme="foo", + code_width=60, + word_wrap=True, + indent_guides=True, + ), + padding=0, + ) + rendered_syntax = render(syntax) + print(repr(rendered_syntax)) + expected = '╭────────────────────────────────────────────────────────────────╮\n│\x1b[1;38;2;24;24;24;48;2;248;248;248m \x1b[0m\x1b[38;2;173;173;173;48;2;248;248;248m 2 \x1b[0m\x1b[3;38;2;64;128;128;48;2;248;248;248m│ \x1b[0m\x1b[3;38;2;186;33;33;48;2;248;248;248m"""Iterate and generate a tuple with a flag for first \x1b[0m\x1b[48;2;248;248;248m \x1b[0m│\n│\x1b[48;2;248;248;248m \x1b[0m\x1b[3;38;2;186;33;33;48;2;248;248;248mand last value."""\x1b[0m\x1b[48;2;248;248;248m \x1b[0m│\n│\x1b[1;38;2;24;24;24;48;2;248;248;248m \x1b[0m\x1b[38;2;173;173;173;48;2;248;248;248m 3 \x1b[0m\x1b[3;38;2;64;128;128;48;2;248;248;248m│ \x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248miter_values\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m \x1b[0m\x1b[38;2;102;102;102;48;2;248;248;248m=\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m \x1b[0m\x1b[38;2;0;128;0;48;2;248;248;248miter\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m(\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248mvalues\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m)\x1b[0m\x1b[48;2;248;248;248m \x1b[0m│\n│\x1b[1;38;2;24;24;24;48;2;248;248;248m \x1b[0m\x1b[38;2;173;173;173;48;2;248;248;248m 4 \x1b[0m\x1b[3;38;2;64;128;128;48;2;248;248;248m│ \x1b[0m\x1b[1;38;2;0;128;0;48;2;248;248;248mtry\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m:\x1b[0m\x1b[48;2;248;248;248m \x1b[0m│\n│\x1b[1;38;2;24;24;24;48;2;248;248;248m \x1b[0m\x1b[38;2;173;173;173;48;2;248;248;248m 5 \x1b[0m\x1b[3;38;2;64;128;128;48;2;248;248;248m│ │ \x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248mprevious_value\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m \x1b[0m\x1b[38;2;102;102;102;48;2;248;248;248m=\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m \x1b[0m\x1b[38;2;0;128;0;48;2;248;248;248mnext\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m(\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248miter_values\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m)\x1b[0m\x1b[48;2;248;248;248m \x1b[0m│\n│\x1b[1;38;2;24;24;24;48;2;248;248;248m \x1b[0m\x1b[38;2;173;173;173;48;2;248;248;248m 6 \x1b[0m\x1b[3;38;2;64;128;128;48;2;248;248;248m│ \x1b[0m\x1b[1;38;2;0;128;0;48;2;248;248;248mexcept\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m \x1b[0m\x1b[1;38;2;210;65;58;48;2;248;248;248mStopIteration\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m:\x1b[0m\x1b[48;2;248;248;248m \x1b[0m│\n│\x1b[1;38;2;24;24;24;48;2;248;248;248m \x1b[0m\x1b[38;2;173;173;173;48;2;248;248;248m 7 \x1b[0m\x1b[3;38;2;64;128;128;48;2;248;248;248m│ │ \x1b[0m\x1b[1;38;2;0;128;0;48;2;248;248;248mreturn\x1b[0m\x1b[48;2;248;248;248m \x1b[0m│\n│\x1b[1;38;2;24;24;24;48;2;248;248;248m \x1b[0m\x1b[38;2;173;173;173;48;2;248;248;248m 8 \x1b[0m\x1b[3;38;2;64;128;128;48;2;248;248;248m│ \x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248mfirst\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m \x1b[0m\x1b[38;2;102;102;102;48;2;248;248;248m=\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m \x1b[0m\x1b[1;38;2;0;128;0;48;2;248;248;248mTrue\x1b[0m\x1b[48;2;248;248;248m \x1b[0m│\n│\x1b[1;38;2;24;24;24;48;2;248;248;248m \x1b[0m\x1b[38;2;173;173;173;48;2;248;248;248m 9 \x1b[0m\x1b[3;38;2;64;128;128;48;2;248;248;248m│ \x1b[0m\x1b[1;38;2;0;128;0;48;2;248;248;248mfor\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m \x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248mvalue\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m \x1b[0m\x1b[1;38;2;170;34;255;48;2;248;248;248min\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m \x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248miter_values\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m:\x1b[0m\x1b[48;2;248;248;248m \x1b[0m│\n│\x1b[1;38;2;24;24;24;48;2;248;248;248m \x1b[0m\x1b[38;2;173;173;173;48;2;248;248;248m10 \x1b[0m\x1b[3;38;2;64;128;128;48;2;248;248;248m│ │ \x1b[0m\x1b[1;38;2;0;128;0;48;2;248;248;248myield\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m \x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248mfirst\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m,\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m \x1b[0m\x1b[1;38;2;0;128;0;48;2;248;248;248mFalse\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m,\x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248m \x1b[0m\x1b[38;2;0;0;0;48;2;248;248;248mprevious_value\x1b[0m\x1b[48;2;248;248;248m \x1b[0m│\n╰────────────────────────────────────────────────────────────────╯\n' + assert rendered_syntax == expected + + def test_pygments_syntax_theme_non_str(): from pygments.style import Style as PygmentsStyle diff --git a/tests/test_text.py b/tests/test_text.py index fc3b1108..5cfc288a 100644 --- a/tests/test_text.py +++ b/tests/test_text.py @@ -598,3 +598,41 @@ def test_align_center(): test = Text("foo") test.align("center", 10) assert test.plain == " foo " + + +def test_detect_indentation(): + test = """\ +foo + bar + """ + assert Text(test).detect_indentation() == 4 + test = """\ +foo + bar + baz + """ + assert Text(test).detect_indentation() == 2 + assert Text("").detect_indentation() == 1 + assert Text(" ").detect_indentation() == 1 + + +def test_indentation_guides(): + test = Text( + """\ +for a in range(10): + print(a) + +foo = [ + 1, + { + 2 + } +] + +""" + ) + result = test.with_indent_guides() + print(result.plain) + print(repr(result.plain)) + expected = "for a in range(10):\n│ print(a)\n\nfoo = [\n│ 1,\n│ {\n│ │ 2\n│ }\n]\n" + assert result.plain == expected \ No newline at end of file diff --git a/tools/stress_test_pretty.py b/tools/stress_test_pretty.py index 8f428ca6..eb27f330 100644 --- a/tools/stress_test_pretty.py +++ b/tools/stress_test_pretty.py @@ -16,4 +16,4 @@ DATA = { } console = Console() for w in range(130): - console.print(Panel(Pretty(DATA), width=w)) + console.print(Panel(Pretty(DATA, indent_guides=True), width=w))