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 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)
