rich/tools/make_terminal_widths.py

86 lines
2.0 KiB
Python

import bisect
import itertools
from operator import itemgetter
import subprocess
from typing import List, Tuple
import sys
from rich.progress import Progress
from wcwidth import wcwidth
progress = Progress()
def make_widths_table() -> List[Tuple[int, int, int]]:
widths: List[Tuple[int, int, int]] = []
make_table_task = progress.add_task("Calculating table...")
cp_widths = (
(codepoint, wcwidth(chr(codepoint)))
for codepoint in range(0, sys.maxunicode + 1)
)
progress.update(make_table_task, total=sys.maxunicode)
for width, codepoints in itertools.groupby(cp_widths, key=itemgetter(1)):
cp_list = list(codepoints)
progress.advance(make_table_task, len(cp_list))
if width == 1:
continue
widths.append((cp_list[0][0], cp_list[-1][0], width))
return widths
def get_cell_size(
widths: List[Tuple[int, int, int]],
codepoint: int,
) -> int:
"""Get the cell size of a character.
Args:
codepoint (int): Codepoint of a character.
Returns:
int: Number of cells (0, 1 or 2) occupied by that character.
"""
idx = bisect.bisect_right(widths, (codepoint, sys.maxunicode + 2))
_start, end, width = widths[idx - 1]
if codepoint <= end:
return width
else:
return 1
def test(widths: List[Tuple[int, int, int]]) -> None:
for codepoint in progress.track(
range(0, sys.maxunicode + 1), description="Testing..."
):
character = chr(codepoint)
width1 = get_cell_size(widths, codepoint)
width2 = wcwidth(character)
if width1 != width2:
print(f"{codepoint}: {width1} != {width2}")
break
def run() -> None:
with progress:
widths = make_widths_table()
test(widths)
table_file = f"""# Auto generated by make_terminal_widths.py
CELL_WIDTHS = {widths!r}
"""
with open("../rich/_cell_widths.py", "wt") as fh:
fh.write(table_file)
subprocess.run("black ../rich/_cell_widths.py", shell=True)
if __name__ == "__main__":
run()