| import re | |
| import unicodedata | |
| import functools | |
| ANSI_ESCAPE_SEQUENCE = re.compile(r"\x1b\[[ -@]*[A-~]") | |
| @functools.cache | |
| def str_width(c: str) -> int: | |
| if ord(c) < 128: | |
| return 1 | |
| w = unicodedata.east_asian_width(c) | |
| if w in ('N', 'Na', 'H', 'A'): | |
| return 1 | |
| return 2 | |
| def wlen(s: str) -> int: | |
| if len(s) == 1: | |
| return str_width(s) | |
| length = sum(str_width(i) for i in s) | |
| # remove lengths of any escape sequences | |
| sequence = ANSI_ESCAPE_SEQUENCE.findall(s) | |
| return length - sum(len(i) for i in sequence) |