diff --git a/witchie/wcstring.py b/witchie/wcstring.py index 9db6ec1..2bac157 100644 --- a/witchie/wcstring.py +++ b/witchie/wcstring.py @@ -55,25 +55,41 @@ def wc_wrap(text: str, length: int) -> Generator[str, None, None]: line_len = 0 words = re.split(r"\s+", text.strip()) + stack = [] # stack to ensure enclosure of style tags for word in words: word_len = wcswidth(word) matches = re.findall(STYLE_TAG_PATTERN, word) for match in matches: - full, _, name = match + full, end, name = match if name in STYLES: word_len -= len(full) + if end != '/': + stack.append(name) + elif len(stack) and name == stack[-1]: + stack.pop() if line_words and line_len + word_len > length: line = " ".join(line_words) - if line_len <= length: + temp_length = length + for style in reversed(stack): + line += '' + temp_length += 3 + len(style) + if line_len <= temp_length: yield line else: - yield from _wc_hard_wrap(line, length) + yield from _wc_hard_wrap(line, temp_length) - line_words = [] line_len = 0 + line_words = [] - line_words.append(word) + if len(line_words) == 0: + styles = '' + for style in stack: + styles += '<' + style + '>' + line_words = [styles + word] + stack = [] + else: + line_words.append(word) line_len += word_len + 1 # add 1 to account for space between words if line_words: