Fix wrapping overflowing lines

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 += '</' + style + '>'
+                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: