From 5e4a6a21bf1f511232c620a927b9906ea295573a Mon Sep 17 00:00:00 2001 From: Kevin Van Brunt Date: Tue, 16 Sep 2025 03:06:20 -0400 Subject: [PATCH] Fixed bug where rich_utils._apply_style_wrapper() could strip a trailing newline. --- cmd2/rich_utils.py | 4 +- tests/test_rich_utils.py | 94 +++++++++++++++++++++++++++++++--------- 2 files changed, 76 insertions(+), 22 deletions(-) diff --git a/cmd2/rich_utils.py b/cmd2/rich_utils.py index 396d96bc3..e5a52db55 100644 --- a/cmd2/rich_utils.py +++ b/cmd2/rich_utils.py @@ -416,9 +416,9 @@ def _apply_style_wrapper(cls: type[Segment], *args: Any, **kwargs: Any) -> Itera styled_segments = list(_orig_segment_apply_style(*args, **kwargs)) newline_segment = cls.line() - # If the final segment is a newline, it will be stripped by Segment.split_lines(). + # If the final segment ends in a newline, that newline will be stripped by Segment.split_lines(). # Save an unstyled newline to restore later. - end_segment = newline_segment if styled_segments and styled_segments[-1].text == "\n" else None + end_segment = newline_segment if styled_segments and styled_segments[-1].text.endswith("\n") else None # Use Segment.split_lines() to separate the styled text from the newlines. # This way the ANSI reset code will appear before any newline. diff --git a/tests/test_rich_utils.py b/tests/test_rich_utils.py index 5f6a07981..dad300bf8 100644 --- a/tests/test_rich_utils.py +++ b/tests/test_rich_utils.py @@ -147,101 +147,106 @@ def test_from_ansi_wrapper() -> None: @pytest.mark.parametrize( # Print with style and verify that everything but newline characters have style. - ('objects', 'expected', 'sep', 'end'), + ('objects', 'sep', 'end', 'expected'), [ # Print nothing - ((), "\n", " ", "\n"), + ((), " ", "\n", "\n"), # Empty string - (("",), "\n", " ", "\n"), + (("",), " ", "\n", "\n"), # Multple empty strings - (("", ""), '\x1b[34;47m \x1b[0m\n', " ", "\n"), + (("", ""), " ", "\n", "\x1b[34;47m \x1b[0m\n"), # Basic string ( ("str_1",), - "\x1b[34;47mstr_1\x1b[0m\n", " ", "\n", + "\x1b[34;47mstr_1\x1b[0m\n", ), # String which ends with newline ( ("str_1\n",), - "\x1b[34;47mstr_1\x1b[0m\n\n", " ", "\n", + "\x1b[34;47mstr_1\x1b[0m\n\n", ), # String which ends with multiple newlines ( ("str_1\n\n",), - "\x1b[34;47mstr_1\x1b[0m\n\n\n", " ", "\n", + "\x1b[34;47mstr_1\x1b[0m\n\n\n", ), # Mutiple lines ( ("str_1\nstr_2",), - "\x1b[34;47mstr_1\x1b[0m\n\x1b[34;47mstr_2\x1b[0m\n", " ", "\n", + "\x1b[34;47mstr_1\x1b[0m\n\x1b[34;47mstr_2\x1b[0m\n", ), # Multiple strings ( ("str_1", "str_2"), - "\x1b[34;47mstr_1 str_2\x1b[0m\n", " ", "\n", + "\x1b[34;47mstr_1 str_2\x1b[0m\n", ), # Multiple strings with newline between them. ( ("str_1\n", "str_2"), - "\x1b[34;47mstr_1\x1b[0m\n\x1b[34;47m str_2\x1b[0m\n", " ", "\n", + "\x1b[34;47mstr_1\x1b[0m\n\x1b[34;47m str_2\x1b[0m\n", ), # Multiple strings and non-space value for sep ( ("str_1", "str_2"), - "\x1b[34;47mstr_1(sep)str_2\x1b[0m\n", "(sep)", "\n", + "\x1b[34;47mstr_1(sep)str_2\x1b[0m\n", ), # Multiple strings and sep is a newline ( ("str_1", "str_2"), - "\x1b[34;47mstr_1\x1b[0m\n\x1b[34;47mstr_2\x1b[0m\n", "\n", "\n", + "\x1b[34;47mstr_1\x1b[0m\n\x1b[34;47mstr_2\x1b[0m\n", ), # Multiple strings and sep has newlines ( ("str_1", "str_2"), - "\x1b[34;47mstr_1(sep1)\x1b[0m\n\x1b[34;47m(sep2)str_2\x1b[0m\n", - "(sep1)\n(sep2)", + "(sep1)\n(sep2)\n", "\n", + ("\x1b[34;47mstr_1(sep1)\x1b[0m\n\x1b[34;47m(sep2)\x1b[0m\n\x1b[34;47mstr_2\x1b[0m\n"), ), # Non-newline value for end. ( ("str_1", "str_2"), - "\x1b[34;47mstr_1(sep1)\x1b[0m\n\x1b[34;47m(sep2)str_2\x1b[0m\x1b[34;47m(end)\x1b[0m", "(sep1)\n(sep2)", "(end)", + "\x1b[34;47mstr_1(sep1)\x1b[0m\n\x1b[34;47m(sep2)str_2\x1b[0m\x1b[34;47m(end)\x1b[0m", ), # end has newlines. ( ("str_1", "str_2"), - "\x1b[34;47mstr_1(sep1)\x1b[0m\n\x1b[34;47m(sep2)str_2\x1b[0m\x1b[34;47m(end1)\x1b[0m\n\x1b[34;47m(end2)\x1b[0m", - "(sep1)\n(sep2)", - "(end1)\n(end2)", + "(sep1)\n(sep2)\n", + "(end1)\n(end2)\n", + ( + "\x1b[34;47mstr_1(sep1)\x1b[0m\n" + "\x1b[34;47m(sep2)\x1b[0m\n" + "\x1b[34;47mstr_2\x1b[0m\x1b[34;47m(end1)\x1b[0m\n" + "\x1b[34;47m(end2)\x1b[0m\n" + ), ), # Empty sep and end values ( ("str_1", "str_2"), - "\x1b[34;47mstr_1str_2\x1b[0m", "", "", + "\x1b[34;47mstr_1str_2\x1b[0m", ), ], ) -def test_apply_style_wrapper(objects: tuple[str], expected: str, sep: str, end: str) -> None: +def test_apply_style_wrapper_soft_wrap(objects: tuple[str], sep: str, end: str, expected: str) -> None: # Check if we are still patching Segment.apply_style(). If this check fails, then Rich # has fixed the bug. Therefore, we can remove this test function and ru._apply_style_wrapper. assert Segment.apply_style.__func__ is ru._apply_style_wrapper.__func__ # type: ignore[attr-defined] @@ -275,3 +280,52 @@ def test_apply_style_wrapper(objects: tuple[str], expected: str, sep: str, end: finally: # Restore the patch Segment.apply_style = ru._apply_style_wrapper # type: ignore[assignment] + + +def test_apply_style_wrapper_word_wrap() -> None: + """ + Test that our patch didn't mess up word wrapping. + Make sure it does not insert styled newlines or apply style to existing newlines. + """ + # Check if we are still patching Segment.apply_style(). If this check fails, then Rich + # has fixed the bug. Therefore, we can remove this test function and ru._apply_style_wrapper. + assert Segment.apply_style.__func__ is ru._apply_style_wrapper.__func__ # type: ignore[attr-defined] + + str1 = "this\nwill word wrap\n" + str2 = "and\nso will this\n" + sep = "(sep1)\n(sep2)\n" + end = "(end1)\n(end2)\n" + style = "blue on white" + + # All newlines should appear outside of ANSI style sequences. + expected = ( + "\x1b[34;47mthis\x1b[0m\n" + "\x1b[34;47mwill word \x1b[0m\n" + "\x1b[34;47mwrap\x1b[0m\n" + "\x1b[34;47m(sep1)\x1b[0m\n" + "\x1b[34;47m(sep2)\x1b[0m\n" + "\x1b[34;47mand\x1b[0m\n" + "\x1b[34;47mso will \x1b[0m\n" + "\x1b[34;47mthis\x1b[0m\n" + "\x1b[34;47m(end1)\x1b[0m\n" + "\x1b[34;47m(end2)\x1b[0m\n" + ) + + # Set a width which will cause word wrapping. + console = Console(force_terminal=True, width=10) + + try: + with console.capture() as capture: + console.print(str1, str2, sep=sep, end=end, style=style, soft_wrap=False) + assert capture.get() == expected + + # Now remove our patch and make sure it produced the same result as unpatched Rich. + Segment.apply_style = ru._orig_segment_apply_style # type: ignore[assignment] + + with console.capture() as capture: + console.print(str1, str2, sep=sep, end=end, style=style, soft_wrap=False) + assert capture.get() == expected + + finally: + # Restore the patch + Segment.apply_style = ru._apply_style_wrapper # type: ignore[assignment]