Fixed bug where rich_utils._apply_style_wrapper() could strip a trailing newline. (#1514)

Author: kmvanbrunt

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.

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]