Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions cmd2/rich_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
94 changes: 74 additions & 20 deletions tests/test_rich_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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]
Loading