1import re 2import unicodedata 3import functools 4 5ANSI_ESCAPE_SEQUENCE = re.compile(r"\x1b\[[ -@]*[A-~]") 6 7 8@functools.cache 9def str_width(c: str) -> int: 10 if ord(c) < 128: 11 return 1 12 w = unicodedata.east_asian_width(c) 13 if w in ('N', 'Na', 'H', 'A'): 14 return 1 15 return 2 16 17 18def wlen(s: str) -> int: 19 if len(s) == 1: 20 return str_width(s) 21 length = sum(str_width(i) for i in s) 22 # remove lengths of any escape sequences 23 sequence = ANSI_ESCAPE_SEQUENCE.findall(s) 24 ctrl_z_cnt = s.count('\x1a') 25 return length - sum(len(i) for i in sequence) + ctrl_z_cnt 26