rich

1
import subprocess
2
from typing import List, Tuple
3
import sys
4

5
from rich.progress import Progress
6

7
from wcwidth import wcwidth
8

9

10
progress = Progress()
11

12

13
def make_widths_table() -> List[Tuple[int, int, int]]:
14
    table: List[Tuple[int, int, int]] = []
15
    append = table.append
16

17
    make_table_task = progress.add_task("Calculating table...")
18

19
    widths = (
20
        (codepoint, wcwidth(chr(codepoint)))
21
        for codepoint in range(0, sys.maxunicode + 1)
22
    )
23

24
    _widths = [(codepoint, width) for codepoint, width in widths if width != 1]
25
    iter_widths = iter(_widths)
26

27
    endpoint, group_cell_size = next(iter_widths)
28
    start_codepoint = end_codepoint = endpoint
29
    for codepoint, cell_size in progress.track(
30
        iter_widths, task_id=make_table_task, total=len(_widths) - 1
31
    ):
32
        if cell_size != group_cell_size or codepoint != end_codepoint + 1:
33
            append((start_codepoint, end_codepoint, group_cell_size))
34
            start_codepoint = end_codepoint = codepoint
35
            group_cell_size = cell_size
36
        else:
37
            end_codepoint = codepoint
38
    append((start_codepoint, end_codepoint, group_cell_size))
39
    return table
40

41

42
def get_cell_size(table: List[Tuple[int, int, int]], character: str) -> int:
43
    codepoint = ord(character)
44
    lower_bound = 0
45
    upper_bound = len(table) - 1
46
    index = (lower_bound + upper_bound) // 2
47
    while True:
48
        start, end, width = table[index]
49
        if codepoint < start:
50
            upper_bound = index - 1
51
        elif codepoint > end:
52
            lower_bound = index + 1
53
        else:
54
            return width
55
        if upper_bound < lower_bound:
56
            break
57
        index = (lower_bound + upper_bound) // 2
58
    return 1
59

60

61
def test(widths_table):
62
    for codepoint in progress.track(
63
        range(0, sys.maxunicode + 1), description="Testing..."
64
    ):
65
        character = chr(codepoint)
66
        width1 = get_cell_size(widths_table, character)
67
        width2 = wcwidth(character)
68
        if width1 != width2:
69
            print(f"{width1} != {width2}")
70
            break
71

72

73
def run():
74
    with progress:
75
        widths_table = make_widths_table()
76
        test(widths_table)
77
    table_file = f"""# Auto generated by make_terminal_widths.py
78

79
CELL_WIDTHS = {widths_table!r}
80

81
"""
82
    with open("../rich/_cell_widths.py", "wt") as fh:
83
        fh.write(table_file)
84

85
    subprocess.run("black ../rich/_cell_widths.py", shell=True)
86

87

88
if __name__ == "__main__":
89
    run()
90