_common.py 3.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. # Source code: https://github.com/hamdanal/rich-argparse
  2. # MIT license: Copyright (c) Ali Hamdan <ali.hamdan.dev@gmail.com>
  3. # for internal use only
  4. from __future__ import annotations
  5. import sys
  6. import rich_argparse._lazy_rich as r
  7. # Default highlight patterns:
  8. # - highlight `text in backquotes` as "syntax"
  9. # - --words-with-dashes outside backticks as "args"
  10. _HIGHLIGHTS = [
  11. r"`(?P<syntax>[^`]*)`|(?:^|\s)(?P<args>-{1,2}[\w]+[\w-]*)",
  12. ]
  13. _windows_console_fixed: bool | None = None
  14. def rich_strip(text: r.Text) -> r.Text:
  15. """Strip leading and trailing whitespace from `rich.text.Text`."""
  16. lstrip_at = len(text.plain) - len(text.plain.lstrip())
  17. if lstrip_at: # rich.Text.lstrip() is not available yet!!
  18. text = text[lstrip_at:]
  19. text.rstrip()
  20. return text
  21. def rich_wrap(console: r.Console, text: r.Text, width: int) -> r.Lines:
  22. """`textwrap.wrap()` equivalent for `rich.text.Text`."""
  23. text = text.copy()
  24. text.expand_tabs(8) # textwrap expands tabs first
  25. whitespace_trans = dict.fromkeys(map(ord, "\t\n\x0b\x0c\r "), ord(" "))
  26. text.plain = text.plain.translate(whitespace_trans)
  27. return text.wrap(console, width)
  28. def rich_fill(console: r.Console, text: r.Text, width: int, indent: r.Text) -> r.Text:
  29. """`textwrap.fill()` equivalent for `rich.text.Text`."""
  30. lines = rich_wrap(console, text, width)
  31. return r.Text("\n").join(indent + line for line in lines)
  32. def _initialize_win_colors() -> bool: # pragma: no cover
  33. global _windows_console_fixed
  34. assert sys.platform == "win32"
  35. if _windows_console_fixed is None:
  36. winver = sys.getwindowsversion() # type: ignore[attr-defined]
  37. if winver.major < 10 or winver.build < 10586:
  38. try:
  39. import colorama
  40. _windows_console_fixed = isinstance(sys.stdout, colorama.ansitowin32.StreamWrapper)
  41. except Exception:
  42. _windows_console_fixed = False
  43. else:
  44. import ctypes
  45. kernel32 = ctypes.windll.kernel32 # type: ignore[attr-defined]
  46. ENABLE_PROCESSED_OUTPUT = 0x1
  47. ENABLE_WRAP_AT_EOL_OUTPUT = 0x2
  48. ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x4
  49. STD_OUTPUT_HANDLE = -11
  50. kernel32.SetConsoleMode(
  51. kernel32.GetStdHandle(STD_OUTPUT_HANDLE),
  52. ENABLE_PROCESSED_OUTPUT
  53. | ENABLE_WRAP_AT_EOL_OUTPUT
  54. | ENABLE_VIRTUAL_TERMINAL_PROCESSING,
  55. )
  56. _windows_console_fixed = True
  57. return _windows_console_fixed
  58. def _fix_legacy_win_text(console: r.Console, text: str) -> str:
  59. # activate legacy Windows console colors if needed (and available) or strip ANSI escape codes
  60. if (
  61. text
  62. and sys.platform == "win32"
  63. and console.legacy_windows
  64. and console.color_system is not None
  65. and not _initialize_win_colors()
  66. ): # pragma: win32 cover
  67. text = "\n".join(r.re_ansi.sub("", line) for line in text.split("\n"))
  68. return text