Cleanup utils dealing with wcwidth strings

tests/test_utils.py

@@ -2,13 +2,74 @@ from toot import utils
 
 
 def test_pad():
-    text = 'Frank Zappa 🎸'
-    padded = utils.pad(text, 14)
-    assert padded == 'Frank Zappa 🎸'
     # guitar symbol will occupy two cells, so padded text should be 1
     # character shorter
-    assert len(padded) == 13
-    # when truncated, … occupies one cell, so we get full length
-    padded = utils.pad(text, 13)
-    assert padded == 'Frank Zappa …'
-    assert len(padded) == 13
+    text = 'Frank Zappa 🎸'
+
+    # Negative values are basically ignored
+    assert utils.pad(text, -100) is text
+
+    # Padding to length smaller than text length does nothing
+    assert utils.pad(text, 11) is text
+    assert utils.pad(text, 12) is text
+    assert utils.pad(text, 13) is text
+    assert utils.pad(text, 14) is text
+
+    assert utils.pad(text, 15) == 'Frank Zappa 🎸 '
+    assert utils.pad(text, 16) == 'Frank Zappa 🎸  '
+    assert utils.pad(text, 17) == 'Frank Zappa 🎸   '
+    assert utils.pad(text, 18) == 'Frank Zappa 🎸    '
+    assert utils.pad(text, 19) == 'Frank Zappa 🎸     '
+    assert utils.pad(text, 20) == 'Frank Zappa 🎸      '
+
+
+def test_trunc():
+    text = 'Frank Zappa 🎸'
+
+    assert utils.trunc(text, 1)  == '…'
+    assert utils.trunc(text, 2)  == 'F…'
+    assert utils.trunc(text, 3)  == 'Fr…'
+    assert utils.trunc(text, 4)  == 'Fra…'
+    assert utils.trunc(text, 5)  == 'Fran…'
+    assert utils.trunc(text, 6)  == 'Frank…'
+    assert utils.trunc(text, 7)  == 'Frank…'
+    assert utils.trunc(text, 8)  == 'Frank Z…'
+    assert utils.trunc(text, 9)  == 'Frank Za…'
+    assert utils.trunc(text, 10) == 'Frank Zap…'
+    assert utils.trunc(text, 11) == 'Frank Zapp…'
+    assert utils.trunc(text, 12) == 'Frank Zappa…'
+    assert utils.trunc(text, 13) == 'Frank Zappa…'
+
+    # Truncating to length larger than text length does nothing
+    assert utils.trunc(text, 14) is text
+    assert utils.trunc(text, 15) is text
+    assert utils.trunc(text, 16) is text
+    assert utils.trunc(text, 17) is text
+    assert utils.trunc(text, 18) is text
+    assert utils.trunc(text, 19) is text
+    assert utils.trunc(text, 20) is text
+
+
+def test_fit_text():
+    text = 'Frank Zappa 🎸'
+
+    assert utils.fit_text(text, 1)  == '…'
+    assert utils.fit_text(text, 2)  == 'F…'
+    assert utils.fit_text(text, 3)  == 'Fr…'
+    assert utils.fit_text(text, 4)  == 'Fra…'
+    assert utils.fit_text(text, 5)  == 'Fran…'
+    assert utils.fit_text(text, 6)  == 'Frank…'
+    assert utils.fit_text(text, 7)  == 'Frank…'
+    assert utils.fit_text(text, 8)  == 'Frank Z…'
+    assert utils.fit_text(text, 9)  == 'Frank Za…'
+    assert utils.fit_text(text, 10) == 'Frank Zap…'
+    assert utils.fit_text(text, 11) == 'Frank Zapp…'
+    assert utils.fit_text(text, 12) == 'Frank Zappa…'
+    assert utils.fit_text(text, 13) == 'Frank Zappa…'
+    assert utils.fit_text(text, 14) == 'Frank Zappa 🎸'
+    assert utils.fit_text(text, 15) == 'Frank Zappa 🎸 '
+    assert utils.fit_text(text, 16) == 'Frank Zappa 🎸  '
+    assert utils.fit_text(text, 17) == 'Frank Zappa 🎸   '
+    assert utils.fit_text(text, 18) == 'Frank Zappa 🎸    '
+    assert utils.fit_text(text, 19) == 'Frank Zappa 🎸     '
+    assert utils.fit_text(text, 20) == 'Frank Zappa 🎸      '

toot/utils.py

@@ -7,7 +7,7 @@ import unicodedata
 import warnings
 
 from bs4 import BeautifulSoup
-from wcwidth import wcswidth
+from wcwidth import wcwidth, wcswidth
 
 from toot.exceptions import ConsoleError
 
@@ -76,21 +76,59 @@ def assert_domain_exists(domain):
         raise ConsoleError("Domain {} not found".format(domain))
 
 
-def trunc(text, length, text_length=None):
-    """Trims text to given length, if trimmed appends ellipsis."""
-    if text_length is None:
-        text_length = len(text)
+def trunc(text, length):
+    """
+    Truncates text to given length, taking into account wide characters.
+
+    If truncated, the last char is replaced by an elipsis.
+    """
+    if length < 1:
+        raise ValueError("length should be 1 or larger")
+
+    # Remove whitespace first so no unneccesary truncation is done.
+    text = text.strip()
+    text_length = wcswidth(text)
+
     if text_length <= length:
         return text
 
-    return text[:length - 1] + '…'
+    # We cannot just remove n characters from the end since we don't know how
+    # wide these characters are and how it will affect text length.
+    # Use wcwidth to determine how many characters need to be truncated.
+    chars_to_truncate = 0
+    trunc_length = 0
+    for char in reversed(text):
+        chars_to_truncate += 1
+        trunc_length += wcwidth(char)
+        if text_length - trunc_length <= length:
+            break
+
+    # Additional char to make room for elipsis
+    n = chars_to_truncate + 1
+    return text[:-n].strip() + '…'
 
 
-def pad(text, length, fill=' '):
+def pad(text, length):
+    """Pads text to given length, taking into account wide characters."""
     text_length = wcswidth(text)
-    text = trunc(text, length, text_length)
-    assert len(text) <= length
-    return text + fill * (length - text_length)
+
+    if text_length < length:
+        return text + ' ' * (length - text_length)
+
+    return text
+
+
+def fit_text(text, length):
+    """Makes text fit the given length by padding or truncating it."""
+    text_length = wcswidth(text)
+
+    if text_length > length:
+        return trunc(text, length)
+
+    if text_length < length:
+        return pad(text, length)
+
+    return text
 
 
 EOF_KEY = "Ctrl-Z" if os.name == 'nt' else "Ctrl-D"