Added a box method and fixed the broken justify method.

2 files changed, 69 insertions, 22 deletions

diff --git a/string_utils.py b/string_utils.py
index c995070..a766ada 100644
--- a/string_utils.py
+++ b/string_utils.py
@@ -1680,6 +1680,20 @@ def replace_all(in_str: str, replace_set: str, replacement: str) -> str:
     return in_str
 
 
+def replace_nth(string: str, source: str, target: str, nth: int):
+    """Replaces the nth occurrance of a substring within a string.
+
+    >>> replace_nth('this is a test', ' ', '-', 3)
+    'this is a-test'
+
+    """
+    where = [m.start() for m in re.finditer(source, string)][nth - 1]
+    before = string[:where]
+    after = string[where:]
+    after = after.replace(source, target, 1)
+    return before + after
+
+
 if __name__ == '__main__':
     import doctest
 
diff --git a/text_utils.py b/text_utils.py
index 564d67e..39b694c 100644
--- a/text_utils.py
+++ b/text_utils.py
@@ -6,6 +6,7 @@
 import contextlib
 import logging
 import math
+import re
 import sys
 from collections import defaultdict
 from dataclasses import dataclass
@@ -137,30 +138,28 @@ def distribute_strings(
     strings: List[str],
     *,
     width: int = 80,
-    alignment: str = "c",
     padding: str = " ",
 ) -> str:
     """
-    Distributes strings into a line with a particular justification.
+    Distributes strings into a line for justified text.
 
     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40)
-    '   this       is         a       test   '
-    >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='l')
-    'this      is        a         test      '
-    >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='r')
-    '      this        is         a      test'
+    '      this      is      a      test     '
 
     """
-    subwidth = math.floor(width / len(strings))
-    retval = ""
-    for string in strings:
-        string = justify_string(string, width=subwidth, alignment=alignment, padding=padding)
-        retval += string
-    while len(retval) > width:
-        retval = retval.replace('  ', ' ', 1)
-    while len(retval) < width:
-        retval = retval.replace(' ', '  ', 1)
-    return retval
+    ret = ' ' + ' '.join(strings) + ' '
+    assert len(ret) < width
+    x = 0
+    while len(ret) < width:
+        spaces = [m.start() for m in re.finditer(r' ([^ ]|$)', ret)]
+        where = spaces[x]
+        before = ret[:where]
+        after = ret[where:]
+        ret = before + padding + after
+        x += 1
+        if x >= len(spaces):
+            x = 0
+    return ret
 
 
 def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") -> str:
@@ -168,15 +167,16 @@ def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") ->
     Justifies a string.
 
     >>> justify_string_by_chunk("This is a test", 40)
-    'This       is              a        test'
+    'This          is          a         test'
     >>> justify_string_by_chunk("This is a test", 20)
-    'This  is    a   test'
+    'This   is   a   test'
 
     """
+    assert len(string) <= width
     padding = padding[0]
     first, *rest, last = string.split()
-    w = width - (len(first) + 1 + len(last) + 1)
-    ret = first + padding + distribute_strings(rest, width=w, padding=padding) + padding + last
+    w = width - (len(first) + len(last))
+    ret = first + distribute_strings(rest, width=w, padding=padding) + last
     return ret
 
 
@@ -192,7 +192,7 @@ def justify_string(
     >>> justify_string('This is another test', width=40, alignment='r')
     '                    This is another test'
     >>> justify_string('This is another test', width=40, alignment='j')
-    'This       is           another     test'
+    'This        is        another       test'
 
     """
     alignment = alignment[0]
@@ -327,6 +327,39 @@ def header(title: str, *, width: int = 80, color: str = ''):
         return ''
 
 
+def box(
+    title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
+):
+    """Draws a box with nice rounded corners.
+
+    >>> box('Title', 'This is text', width=30)
+    ╭────────────────────────────╮
+    │            Title           │
+    │                            │
+    │ This is text               │
+    ╰────────────────────────────╯
+
+    """
+    assert width > 4
+    if color == '':
+        rset = ''
+    else:
+        rset = reset()
+    w = width - 2
+    print(color + '╭' + '─' * w + '╮' + rset)
+    if title is not None:
+        print(
+            color + '│' + rset + justify_string(title, width=w, alignment='c') + color + '│' + rset
+        )
+        print(color + '│' + ' ' * w + '│' + rset)
+    if text is not None:
+        for line in justify_text(text, width=w - 2, alignment='l').split('\n'):
+            tw = len(line)
+            assert tw < w
+            print(color + '│ ' + rset + line + ' ' * (w - tw - 2) + color + ' │' + rset)
+    print(color + '╰' + '─' * w + '╯' + rset)
+
+
 if __name__ == '__main__':
     import doctest