-
Notifications
You must be signed in to change notification settings - Fork 128
Expand file tree
/
Copy pathrich_utils.py
More file actions
540 lines (435 loc) · 18.9 KB
/
rich_utils.py
File metadata and controls
540 lines (435 loc) · 18.9 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
"""Provides common utilities to support Rich in cmd2-based applications."""
import re
import threading
from collections.abc import Mapping
from enum import Enum
from typing import (
IO,
Any,
TypedDict,
)
from rich.console import (
Console,
ConsoleRenderable,
JustifyMethod,
OverflowMethod,
RenderableType,
)
from rich.padding import Padding
from rich.pretty import is_expandable
from rich.protocol import rich_cast
from rich.style import StyleType
from rich.table import (
Column,
Table,
)
from rich.text import Text
from rich.theme import Theme
from rich_argparse import RichHelpFormatter
from .styles import DEFAULT_CMD2_STYLES
# Matches ANSI SGR (Select Graphic Rendition) sequences for text styling.
# \x1b[ - the CSI (Control Sequence Introducer)
# [0-9;]* - zero or more digits or semicolons (parameters for the style)
# m - the SGR final character
ANSI_STYLE_SEQUENCE_RE = re.compile(r"\x1b\[[0-9;]*m")
class AllowStyle(Enum):
"""Values for ``cmd2.rich_utils.ALLOW_STYLE``."""
ALWAYS = "Always" # Always output ANSI style sequences
NEVER = "Never" # Remove ANSI style sequences from all output
TERMINAL = "Terminal" # Remove ANSI style sequences if the output is not going to the terminal
def __str__(self) -> str:
"""Return value instead of enum name for printing in cmd2's set command."""
return str(self.value)
def __repr__(self) -> str:
"""Return quoted value instead of enum description for printing in cmd2's set command."""
return repr(self.value)
# Controls when ANSI style sequences are allowed in output
ALLOW_STYLE = AllowStyle.TERMINAL
def _create_default_theme() -> Theme:
"""Create a default theme for the application.
This theme combines the default styles from cmd2, rich-argparse, and Rich.
"""
app_styles = DEFAULT_CMD2_STYLES.copy()
app_styles.update(RichHelpFormatter.styles.copy())
return Theme(app_styles, inherit=True)
def set_theme(styles: Mapping[str, StyleType] | None = None) -> None:
"""Set the Rich theme used by cmd2.
Call set_theme() with no arguments to reset to the default theme.
This will clear any custom styles that were previously applied.
:param styles: optional mapping of style names to styles
"""
global APP_THEME # noqa: PLW0603
# Start with a fresh copy of the default styles.
app_styles: dict[str, StyleType] = {}
app_styles.update(_create_default_theme().styles)
# Incorporate custom styles.
if styles is not None:
app_styles.update(styles)
APP_THEME = Theme(app_styles)
# Synchronize rich-argparse styles with the main application theme.
for name in RichHelpFormatter.styles.keys() & APP_THEME.styles.keys():
RichHelpFormatter.styles[name] = APP_THEME.styles[name]
# The application-wide theme. You can change it with set_theme().
APP_THEME = _create_default_theme()
class RichPrintKwargs(TypedDict, total=False):
"""Keyword arguments that can be passed to rich.console.Console.print() via cmd2's print methods.
See Rich's Console.print() documentation for full details on these parameters.
https://rich.readthedocs.io/en/stable/reference/console.html#rich.console.Console.print
Note: All fields are optional (total=False). If a key is not present in the
dictionary, Rich's default behavior for that argument will apply.
"""
justify: JustifyMethod | None
overflow: OverflowMethod | None
no_wrap: bool | None
width: int | None
height: int | None
crop: bool
new_line_start: bool
class Cmd2BaseConsole(Console):
"""Base class for all cmd2 Rich consoles.
This class handles the core logic for managing Rich behavior based on
cmd2's global settings, such as `ALLOW_STYLE` and `APP_THEME`.
"""
def __init__(
self,
*,
file: IO[str] | None = None,
**kwargs: Any,
) -> None:
"""Cmd2BaseConsole initializer.
:param file: optional file object where the console should write to.
Defaults to sys.stdout.
:param kwargs: keyword arguments passed to the parent Console class.
:raises TypeError: if disallowed keyword argument is passed in.
"""
# Don't allow force_terminal or force_interactive to be passed in, as their
# behavior is controlled by the ALLOW_STYLE setting.
if "force_terminal" in kwargs:
raise TypeError(
"Passing 'force_terminal' is not allowed. Its behavior is controlled by the 'ALLOW_STYLE' setting."
)
if "force_interactive" in kwargs:
raise TypeError(
"Passing 'force_interactive' is not allowed. Its behavior is controlled by the 'ALLOW_STYLE' setting."
)
# Don't allow a theme to be passed in, as it is controlled by the global APP_THEME.
# Use cmd2.rich_utils.set_theme() to set the global theme or use a temporary
# theme with console.use_theme().
if "theme" in kwargs:
raise TypeError(
"Passing 'theme' is not allowed. Its behavior is controlled by the global APP_THEME and set_theme()."
)
force_terminal: bool | None = None
force_interactive: bool | None = None
if ALLOW_STYLE == AllowStyle.ALWAYS:
force_terminal = True
# Turn off interactive mode if dest is not actually a terminal which supports it
tmp_console = Console(file=file)
force_interactive = tmp_console.is_interactive
elif ALLOW_STYLE == AllowStyle.NEVER:
force_terminal = False
super().__init__(
file=file,
force_terminal=force_terminal,
force_interactive=force_interactive,
theme=APP_THEME,
**kwargs,
)
self._thread_local = threading.local()
def on_broken_pipe(self) -> None:
"""Override which raises BrokenPipeError instead of SystemExit."""
self.quiet = True
raise BrokenPipeError
def render_str(
self,
text: str,
highlight: bool | None = None,
markup: bool | None = None,
emoji: bool | None = None,
**kwargs: Any,
) -> Text:
"""Override to ensure formatting overrides passed to print() and log() are respected."""
if emoji is None:
emoji = getattr(self._thread_local, "emoji", None)
if markup is None:
markup = getattr(self._thread_local, "markup", None)
if highlight is None:
highlight = getattr(self._thread_local, "highlight", None)
return super().render_str(text, highlight=highlight, markup=markup, emoji=emoji, **kwargs)
def print(
self,
*objects: Any,
sep: str = " ",
end: str = "\n",
style: StyleType | None = None,
justify: JustifyMethod | None = None,
overflow: OverflowMethod | None = None,
no_wrap: bool | None = None,
emoji: bool | None = None,
markup: bool | None = None,
highlight: bool | None = None,
width: int | None = None,
height: int | None = None,
crop: bool = True,
soft_wrap: bool | None = None,
new_line_start: bool = False,
) -> None:
"""Override to support ANSI sequences and address a bug in Rich.
This method calls [cmd2.rich_utils.prepare_objects_for_rendering][] on the
objects being printed. This ensures that strings containing ANSI style
sequences are converted to Rich Text objects, so that Rich can correctly
calculate their display width.
Additionally, it works around a bug in Rich where complex renderables
(like Table and Rule) may not receive formatting settings passed to print().
By temporarily injecting these settings into thread-local storage, we ensure
that all internal rendering calls within the print() operation respect the
requested overrides.
There is an issue on Rich to fix the latter:
https://github.com/Textualize/rich/issues/4028
"""
prepared_objects = prepare_objects_for_rendering(*objects)
# Inject overrides into thread-local storage
self._thread_local.emoji = emoji
self._thread_local.markup = markup
self._thread_local.highlight = highlight
try:
super().print(
*prepared_objects,
sep=sep,
end=end,
style=style,
justify=justify,
overflow=overflow,
no_wrap=no_wrap,
emoji=emoji,
markup=markup,
highlight=highlight,
width=width,
height=height,
crop=crop,
soft_wrap=soft_wrap,
new_line_start=new_line_start,
)
finally:
# Clear overrides from thread-local storage
self._thread_local.emoji = None
self._thread_local.markup = None
self._thread_local.highlight = None
def log(
self,
*objects: Any,
sep: str = " ",
end: str = "\n",
style: StyleType | None = None,
justify: JustifyMethod | None = None,
emoji: bool | None = None,
markup: bool | None = None,
highlight: bool | None = None,
log_locals: bool = False,
_stack_offset: int = 1,
) -> None:
"""Override to support ANSI sequences and address a bug in Rich.
This method calls [cmd2.rich_utils.prepare_objects_for_rendering][] on the
objects being logged. This ensures that strings containing ANSI style
sequences are converted to Rich Text objects, so that Rich can correctly
calculate their display width.
Additionally, it works around a bug in Rich where complex renderables
(like Table and Rule) may not receive formatting settings passed to log().
By temporarily injecting these settings into thread-local storage, we ensure
that all internal rendering calls within the log() operation respect the
requested overrides.
There is an issue on Rich to fix the latter:
https://github.com/Textualize/rich/issues/4028
"""
prepared_objects = prepare_objects_for_rendering(*objects)
# Inject overrides into thread-local storage
self._thread_local.emoji = emoji
self._thread_local.markup = markup
self._thread_local.highlight = highlight
try:
# Increment _stack_offset because we added this wrapper frame
super().log(
*prepared_objects,
sep=sep,
end=end,
style=style,
justify=justify,
emoji=emoji,
markup=markup,
highlight=highlight,
log_locals=log_locals,
_stack_offset=_stack_offset + 1,
)
finally:
# Clear overrides from thread-local storage
self._thread_local.emoji = None
self._thread_local.markup = None
self._thread_local.highlight = None
class Cmd2GeneralConsole(Cmd2BaseConsole):
"""Rich console for general-purpose printing.
It enables soft wrap and disables Rich's automatic detection for markup,
emoji, and highlighting. These defaults can be overridden in calls to the
console's or cmd2's print methods.
"""
def __init__(self, *, file: IO[str] | None = None) -> None:
"""Cmd2GeneralConsole initializer.
:param file: optional file object where the console should write to.
Defaults to sys.stdout.
"""
super().__init__(
file=file,
soft_wrap=True,
markup=False,
emoji=False,
highlight=False,
)
class Cmd2RichArgparseConsole(Cmd2BaseConsole):
"""Rich console for rich-argparse output.
Ensures long lines in help text are not truncated by disabling soft_wrap,
which conflicts with rich-argparse's explicit no_wrap and overflow settings.
Since this console is used to print error messages which may not be intended
for Rich formatting, it disables Rich's automatic detection for markup, emoji,
and highlighting. Because rich-argparse does markup and highlighting without
involving the console, disabling these settings does not affect the library's
internal functionality.
"""
def __init__(self, *, file: IO[str] | None = None) -> None:
"""Cmd2RichArgparseConsole initializer.
:param file: optional file object where the console should write to.
Defaults to sys.stdout.
"""
super().__init__(
file=file,
soft_wrap=False,
markup=False,
emoji=False,
highlight=False,
)
class Cmd2ExceptionConsole(Cmd2BaseConsole):
"""Rich console for printing exceptions and Rich Tracebacks.
Ensures that output is always word-wrapped for readability and disables
Rich's automatic detection for markup, emoji, and highlighting to prevent
interference with raw error data.
"""
def __init__(self, *, file: IO[str] | None = None) -> None:
"""Cmd2ExceptionConsole initializer.
:param file: optional file object where the console should write to.
Defaults to sys.stdout.
"""
super().__init__(
file=file,
soft_wrap=False,
markup=False,
emoji=False,
highlight=False,
)
def console_width() -> int:
"""Return the width of the console."""
return Console().width
def rich_text_to_string(text: Text) -> str:
"""Convert a Rich Text object to a string.
This function's purpose is to render a Rich Text object, including any styles (e.g., color, bold),
to a plain Python string with ANSI style sequences. It differs from `text.plain`, which strips
all formatting.
:param text: the text object to convert
:return: the resulting string with ANSI styles preserved.
"""
console = Console(
force_terminal=True,
soft_wrap=True,
no_color=False,
theme=APP_THEME,
)
with console.capture() as capture:
console.print(text, end="")
return capture.get()
def indent(renderable: RenderableType, level: int) -> Padding:
"""Indent a Rich renderable.
When soft-wrapping is enabled, a Rich console is unable to properly print a
Padding object of indented text, as it truncates long strings instead of wrapping
them. This function provides a workaround for this issue, ensuring that indented
text is printed correctly regardless of the soft-wrap setting.
For non-text objects, this function merely serves as a convenience
wrapper around Padding.indent().
:param renderable: a Rich renderable to indent.
:param level: number of characters to indent.
:return: a Padding object containing the indented content.
"""
if isinstance(renderable, (str, Text)):
# Wrap text in a grid to handle the wrapping.
text_grid = Table.grid(Column(overflow="fold"))
text_grid.add_row(renderable)
renderable = text_grid
return Padding.indent(renderable, level)
def prepare_objects_for_rendering(*objects: Any) -> tuple[Any, ...]:
"""Prepare a tuple of objects for printing by Rich's Console.print().
This function processes objects to ensure they are rendered correctly by Rich.
It inspects each object and, if its string representation contains ANSI style
sequences, it converts the object to a Rich Text object. This ensures Rich can
properly parse the non-printing codes for accurate display width calculation.
Objects that already implement the Rich console protocol or are expandable
by its pretty printer are left untouched, as they can be handled directly by
Rich's native renderers.
:param objects: objects to prepare
:return: a tuple containing the processed objects.
"""
object_list = list(objects)
for i, obj in enumerate(object_list):
# Resolve the object's final renderable form, including those
# with a __rich__ method that might return a string.
renderable = rich_cast(obj)
# No preprocessing is needed for Rich-compatible or expandable objects.
if isinstance(renderable, ConsoleRenderable) or is_expandable(renderable):
continue
# Check for ANSI style sequences in its string representation.
renderable_as_str = str(renderable)
if ANSI_STYLE_SEQUENCE_RE.search(renderable_as_str):
object_list[i] = Text.from_ansi(renderable_as_str)
return tuple(object_list)
###################################################################################
# Rich Library Monkey Patches
#
# These patches fix specific bugs in the Rich library. They are conditional and
# will only be applied if the bug is detected. When the bugs are fixed in a
# future Rich release, these patches and their corresponding tests should be
# removed.
###################################################################################
###################################################################################
# Text.from_ansi() monkey patch
###################################################################################
# Save original Text.from_ansi() so we can call it in our wrapper
_orig_text_from_ansi = Text.from_ansi
@classmethod # type: ignore[misc]
def _from_ansi_wrapper(cls: type[Text], text: str, *args: Any, **kwargs: Any) -> Text: # noqa: ARG001
r"""Wrap Text.from_ansi() to fix its trailing newline bug.
This wrapper handles an issue where Text.from_ansi() removes the
trailing line break from a string (e.g. "Hello\n" becomes "Hello").
There is currently a pull request on Rich to fix this.
https://github.com/Textualize/rich/pull/3793
"""
result = _orig_text_from_ansi(text, *args, **kwargs)
# If the original string ends with a recognized line break character,
# then restore the missing newline. We use "\n" because Text.from_ansi()
# converts all line breaks into newlines.
# Source: https://docs.python.org/3/library/stdtypes.html#str.splitlines
line_break_chars = {
"\n", # Line Feed
"\r", # Carriage Return
"\v", # Vertical Tab
"\f", # Form Feed
"\x1c", # File Separator
"\x1d", # Group Separator
"\x1e", # Record Separator
"\x85", # Next Line (NEL)
"\u2028", # Line Separator
"\u2029", # Paragraph Separator
}
if text and text[-1] in line_break_chars:
result.append("\n")
return result
def _from_ansi_has_newline_bug() -> bool:
"""Check if Test.from_ansi() strips the trailing line break from a string."""
return Text.from_ansi("\n") == Text.from_ansi("")
# Only apply the monkey patch if the bug is present
if _from_ansi_has_newline_bug():
Text.from_ansi = _from_ansi_wrapper # type: ignore[assignment]