cells.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174
  1. from __future__ import annotations
  2. from functools import lru_cache
  3. from typing import Callable
  4. from ._cell_widths import CELL_WIDTHS
  5. # Ranges of unicode ordinals that produce a 1-cell wide character
  6. # This is non-exhaustive, but covers most common Western characters
  7. _SINGLE_CELL_UNICODE_RANGES: list[tuple[int, int]] = [
  8. (0x20, 0x7E), # Latin (excluding non-printable)
  9. (0xA0, 0xAC),
  10. (0xAE, 0x002FF),
  11. (0x00370, 0x00482), # Greek / Cyrillic
  12. (0x02500, 0x025FC), # Box drawing, box elements, geometric shapes
  13. (0x02800, 0x028FF), # Braille
  14. ]
  15. # A set of characters that are a single cell wide
  16. _SINGLE_CELLS = frozenset(
  17. [
  18. character
  19. for _start, _end in _SINGLE_CELL_UNICODE_RANGES
  20. for character in map(chr, range(_start, _end + 1))
  21. ]
  22. )
  23. # When called with a string this will return True if all
  24. # characters are single-cell, otherwise False
  25. _is_single_cell_widths: Callable[[str], bool] = _SINGLE_CELLS.issuperset
  26. @lru_cache(4096)
  27. def cached_cell_len(text: str) -> int:
  28. """Get the number of cells required to display text.
  29. This method always caches, which may use up a lot of memory. It is recommended to use
  30. `cell_len` over this method.
  31. Args:
  32. text (str): Text to display.
  33. Returns:
  34. int: Get the number of cells required to display text.
  35. """
  36. if _is_single_cell_widths(text):
  37. return len(text)
  38. return sum(map(get_character_cell_size, text))
  39. def cell_len(text: str, _cell_len: Callable[[str], int] = cached_cell_len) -> int:
  40. """Get the number of cells required to display text.
  41. Args:
  42. text (str): Text to display.
  43. Returns:
  44. int: Get the number of cells required to display text.
  45. """
  46. if len(text) < 512:
  47. return _cell_len(text)
  48. if _is_single_cell_widths(text):
  49. return len(text)
  50. return sum(map(get_character_cell_size, text))
  51. @lru_cache(maxsize=4096)
  52. def get_character_cell_size(character: str) -> int:
  53. """Get the cell size of a character.
  54. Args:
  55. character (str): A single character.
  56. Returns:
  57. int: Number of cells (0, 1 or 2) occupied by that character.
  58. """
  59. codepoint = ord(character)
  60. _table = CELL_WIDTHS
  61. lower_bound = 0
  62. upper_bound = len(_table) - 1
  63. index = (lower_bound + upper_bound) // 2
  64. while True:
  65. start, end, width = _table[index]
  66. if codepoint < start:
  67. upper_bound = index - 1
  68. elif codepoint > end:
  69. lower_bound = index + 1
  70. else:
  71. return 0 if width == -1 else width
  72. if upper_bound < lower_bound:
  73. break
  74. index = (lower_bound + upper_bound) // 2
  75. return 1
  76. def set_cell_size(text: str, total: int) -> str:
  77. """Set the length of a string to fit within given number of cells."""
  78. if _is_single_cell_widths(text):
  79. size = len(text)
  80. if size < total:
  81. return text + " " * (total - size)
  82. return text[:total]
  83. if total <= 0:
  84. return ""
  85. cell_size = cell_len(text)
  86. if cell_size == total:
  87. return text
  88. if cell_size < total:
  89. return text + " " * (total - cell_size)
  90. start = 0
  91. end = len(text)
  92. # Binary search until we find the right size
  93. while True:
  94. pos = (start + end) // 2
  95. before = text[: pos + 1]
  96. before_len = cell_len(before)
  97. if before_len == total + 1 and cell_len(before[-1]) == 2:
  98. return before[:-1] + " "
  99. if before_len == total:
  100. return before
  101. if before_len > total:
  102. end = pos
  103. else:
  104. start = pos
  105. def chop_cells(
  106. text: str,
  107. width: int,
  108. ) -> list[str]:
  109. """Split text into lines such that each line fits within the available (cell) width.
  110. Args:
  111. text: The text to fold such that it fits in the given width.
  112. width: The width available (number of cells).
  113. Returns:
  114. A list of strings such that each string in the list has cell width
  115. less than or equal to the available width.
  116. """
  117. _get_character_cell_size = get_character_cell_size
  118. lines: list[list[str]] = [[]]
  119. append_new_line = lines.append
  120. append_to_last_line = lines[-1].append
  121. total_width = 0
  122. for character in text:
  123. cell_width = _get_character_cell_size(character)
  124. char_doesnt_fit = total_width + cell_width > width
  125. if char_doesnt_fit:
  126. append_new_line([character])
  127. append_to_last_line = lines[-1].append
  128. total_width = cell_width
  129. else:
  130. append_to_last_line(character)
  131. total_width += cell_width
  132. return ["".join(line) for line in lines]
  133. if __name__ == "__main__": # pragma: no cover
  134. print(get_character_cell_size("😽"))
  135. for line in chop_cells("""这是对亚洲语言支持的测试。面对模棱两可的想法,拒绝猜测的诱惑。""", 8):
  136. print(line)
  137. for n in range(80, 1, -1):
  138. print(set_cell_size("""这是对亚洲语言支持的测试。面对模棱两可的想法,拒绝猜测的诱惑。""", n) + "|")
  139. print("x" * n)