_textwrap.py 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188
  1. from __future__ import annotations
  2. import collections.abc as cabc
  3. import textwrap
  4. from contextlib import contextmanager
  5. from ._compat import _ansi_re
  6. from ._compat import term_len
  7. def _truncate_visible(text: str, n: int) -> str:
  8. """Return the longest prefix of ``text`` containing at most ``n`` visible
  9. characters.
  10. ANSI escape sequences inside the prefix are kept intact and do not count
  11. toward the visible width. A cut is never placed inside an escape sequence.
  12. """
  13. if n <= 0:
  14. return ""
  15. visible = 0
  16. i = 0
  17. cut = 0
  18. end = len(text)
  19. while i < end:
  20. m = _ansi_re.match(text, i)
  21. if m is not None:
  22. i = m.end()
  23. continue
  24. visible += 1
  25. i += 1
  26. cut = i
  27. if visible >= n:
  28. break
  29. return text[:cut]
  30. class TextWrapper(textwrap.TextWrapper):
  31. """``textwrap.TextWrapper`` variant that measures widths by visible
  32. character count.
  33. ANSI escape sequences embedded in chunks, indents, or the placeholder are
  34. excluded from the width budget. Without this, styled help text (a styled
  35. ``Usage:`` prefix, a colorized option name, ...) would be wrapped earlier
  36. than its visible length warrants and tokens would split mid-word.
  37. """
  38. def _handle_long_word(
  39. self,
  40. reversed_chunks: list[str],
  41. cur_line: list[str],
  42. cur_len: int,
  43. width: int,
  44. ) -> None:
  45. space_left = max(width - cur_len, 1)
  46. if self.break_long_words:
  47. last = reversed_chunks[-1]
  48. cut = _truncate_visible(last, space_left)
  49. res = last[len(cut) :]
  50. cur_line.append(cut)
  51. reversed_chunks[-1] = res
  52. elif not cur_line:
  53. cur_line.append(reversed_chunks.pop())
  54. def _wrap_chunks(self, chunks: list[str]) -> list[str]:
  55. """Wrap chunks counting widths in visible characters.
  56. Mirrors the algorithm of :meth:`textwrap.TextWrapper._wrap_chunks`
  57. with every width measurement routed through
  58. :func:`click._compat.term_len` instead of :func:`len`, so ANSI escape
  59. bytes in chunks, indents, or the placeholder do not inflate the count.
  60. .. seealso::
  61. :class:`textwrap.TextWrapper` in the Python standard library documentation:
  62. https://docs.python.org/3/library/textwrap.html#textwrap.TextWrapper
  63. Reference implementation in CPython:
  64. https://github.com/python/cpython/blob/main/Lib/textwrap.py
  65. """
  66. lines: list[str] = []
  67. if self.width <= 0:
  68. raise ValueError(f"invalid width {self.width!r} (must be > 0)")
  69. if self.max_lines is not None:
  70. if self.max_lines > 1:
  71. indent = self.subsequent_indent
  72. else:
  73. indent = self.initial_indent
  74. if term_len(indent) + term_len(self.placeholder.lstrip()) > self.width:
  75. raise ValueError("placeholder too large for max width")
  76. chunks.reverse()
  77. while chunks:
  78. cur_line: list[str] = []
  79. cur_len = 0
  80. if lines:
  81. indent = self.subsequent_indent
  82. else:
  83. indent = self.initial_indent
  84. width = self.width - term_len(indent)
  85. if self.drop_whitespace and chunks[-1].strip() == "" and lines:
  86. del chunks[-1]
  87. while chunks:
  88. n = term_len(chunks[-1])
  89. if cur_len + n <= width:
  90. cur_line.append(chunks.pop())
  91. cur_len += n
  92. else:
  93. break
  94. if chunks and term_len(chunks[-1]) > width:
  95. self._handle_long_word(chunks, cur_line, cur_len, width)
  96. cur_len = sum(map(term_len, cur_line))
  97. if self.drop_whitespace and cur_line and cur_line[-1].strip() == "":
  98. cur_len -= term_len(cur_line[-1])
  99. del cur_line[-1]
  100. if cur_line:
  101. if (
  102. self.max_lines is None
  103. or len(lines) + 1 < self.max_lines
  104. or (
  105. not chunks
  106. or self.drop_whitespace
  107. and len(chunks) == 1
  108. and not chunks[0].strip()
  109. )
  110. and cur_len <= width
  111. ):
  112. lines.append(indent + "".join(cur_line))
  113. else:
  114. while cur_line:
  115. if (
  116. cur_line[-1].strip()
  117. and cur_len + term_len(self.placeholder) <= width
  118. ):
  119. cur_line.append(self.placeholder)
  120. lines.append(indent + "".join(cur_line))
  121. break
  122. cur_len -= term_len(cur_line[-1])
  123. del cur_line[-1]
  124. else:
  125. if lines:
  126. prev_line = lines[-1].rstrip()
  127. if (
  128. term_len(prev_line) + term_len(self.placeholder)
  129. <= self.width
  130. ):
  131. lines[-1] = prev_line + self.placeholder
  132. break
  133. lines.append(indent + self.placeholder.lstrip())
  134. break
  135. return lines
  136. @contextmanager
  137. def extra_indent(self, indent: str) -> cabc.Iterator[None]:
  138. old_initial_indent = self.initial_indent
  139. old_subsequent_indent = self.subsequent_indent
  140. self.initial_indent += indent
  141. self.subsequent_indent += indent
  142. try:
  143. yield
  144. finally:
  145. self.initial_indent = old_initial_indent
  146. self.subsequent_indent = old_subsequent_indent
  147. def indent_only(self, text: str) -> str:
  148. rv = []
  149. for idx, line in enumerate(text.splitlines()):
  150. indent = self.initial_indent
  151. if idx > 0:
  152. indent = self.subsequent_indent
  153. rv.append(f"{indent}{line}")
  154. return "\n".join(rv)