cells.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167
  1. from __future__ import annotations
  2. import re
  3. from functools import lru_cache
  4. from typing import Callable
  5. from ._cell_widths import CELL_WIDTHS
  6. # Regex to match sequence of the most common character ranges
  7. _is_single_cell_widths = re.compile("^[\u0020-\u006f\u00a0\u02ff\u0370-\u0482]*$").match
  8. @lru_cache(4096)
  9. def cached_cell_len(text: str) -> int:
  10. """Get the number of cells required to display text.
  11. This method always caches, which may use up a lot of memory. It is recommended to use
  12. `cell_len` over this method.
  13. Args:
  14. text (str): Text to display.
  15. Returns:
  16. int: Get the number of cells required to display text.
  17. """
  18. _get_size = get_character_cell_size
  19. total_size = sum(_get_size(character) for character in text)
  20. return total_size
  21. def cell_len(text: str, _cell_len: Callable[[str], int] = cached_cell_len) -> int:
  22. """Get the number of cells required to display text.
  23. Args:
  24. text (str): Text to display.
  25. Returns:
  26. int: Get the number of cells required to display text.
  27. """
  28. if len(text) < 512:
  29. return _cell_len(text)
  30. _get_size = get_character_cell_size
  31. total_size = sum(_get_size(character) for character in text)
  32. return total_size
  33. @lru_cache(maxsize=4096)
  34. def get_character_cell_size(character: str) -> int:
  35. """Get the cell size of a character.
  36. Args:
  37. character (str): A single character.
  38. Returns:
  39. int: Number of cells (0, 1 or 2) occupied by that character.
  40. """
  41. return _get_codepoint_cell_size(ord(character))
  42. @lru_cache(maxsize=4096)
  43. def _get_codepoint_cell_size(codepoint: int) -> int:
  44. """Get the cell size of a character.
  45. Args:
  46. codepoint (int): Codepoint of a character.
  47. Returns:
  48. int: Number of cells (0, 1 or 2) occupied by that character.
  49. """
  50. _table = CELL_WIDTHS
  51. lower_bound = 0
  52. upper_bound = len(_table) - 1
  53. index = (lower_bound + upper_bound) // 2
  54. while True:
  55. start, end, width = _table[index]
  56. if codepoint < start:
  57. upper_bound = index - 1
  58. elif codepoint > end:
  59. lower_bound = index + 1
  60. else:
  61. return 0 if width == -1 else width
  62. if upper_bound < lower_bound:
  63. break
  64. index = (lower_bound + upper_bound) // 2
  65. return 1
  66. def set_cell_size(text: str, total: int) -> str:
  67. """Set the length of a string to fit within given number of cells."""
  68. if _is_single_cell_widths(text):
  69. size = len(text)
  70. if size < total:
  71. return text + " " * (total - size)
  72. return text[:total]
  73. if total <= 0:
  74. return ""
  75. cell_size = cell_len(text)
  76. if cell_size == total:
  77. return text
  78. if cell_size < total:
  79. return text + " " * (total - cell_size)
  80. start = 0
  81. end = len(text)
  82. # Binary search until we find the right size
  83. while True:
  84. pos = (start + end) // 2
  85. before = text[: pos + 1]
  86. before_len = cell_len(before)
  87. if before_len == total + 1 and cell_len(before[-1]) == 2:
  88. return before[:-1] + " "
  89. if before_len == total:
  90. return before
  91. if before_len > total:
  92. end = pos
  93. else:
  94. start = pos
  95. def chop_cells(
  96. text: str,
  97. width: int,
  98. ) -> list[str]:
  99. """Split text into lines such that each line fits within the available (cell) width.
  100. Args:
  101. text: The text to fold such that it fits in the given width.
  102. width: The width available (number of cells).
  103. Returns:
  104. A list of strings such that each string in the list has cell width
  105. less than or equal to the available width.
  106. """
  107. _get_character_cell_size = get_character_cell_size
  108. lines: list[list[str]] = [[]]
  109. append_new_line = lines.append
  110. append_to_last_line = lines[-1].append
  111. total_width = 0
  112. for character in text:
  113. cell_width = _get_character_cell_size(character)
  114. char_doesnt_fit = total_width + cell_width > width
  115. if char_doesnt_fit:
  116. append_new_line([character])
  117. append_to_last_line = lines[-1].append
  118. total_width = cell_width
  119. else:
  120. append_to_last_line(character)
  121. total_width += cell_width
  122. return ["".join(line) for line in lines]
  123. if __name__ == "__main__": # pragma: no cover
  124. print(get_character_cell_size("😽"))
  125. for line in chop_cells("""这是对亚洲语言支持的测试。面对模棱两可的想法,拒绝猜测的诱惑。""", 8):
  126. print(line)
  127. for n in range(80, 1, -1):
  128. print(set_cell_size("""这是对亚洲语言支持的测试。面对模棱两可的想法,拒绝猜测的诱惑。""", n) + "|")
  129. print("x" * n)