Improve text wrapping; add indent option.

1 files changed, 15 insertions, 8 deletions

diff --git a/text_utils.py b/text_utils.py
index cc1269c..17ac644 100644
--- a/text_utils.py
+++ b/text_utils.py
@@ -217,17 +217,21 @@ def justify_string(
     return string
 
 
-def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str:
+def justify_text(text: str, *, width: int = 80, alignment: str = "c", indent_by: int = 0) -> str:
     """
-    Justifies text.
+    Justifies text optionally with initial indentation.
 
     >>> justify_text('This is a test of the emergency broadcast system.  This is only a test.',
     ...              width=40, alignment='j')  #doctest: +NORMALIZE_WHITESPACE
     'This  is    a  test  of   the  emergency\\nbroadcast system. This is only a test.'
 
     """
-    retval = ""
-    line = ""
+    retval = ''
+    indent = ''
+    if indent_by > 0:
+        indent += ' ' * indent_by
+    line = indent
+
     for word in text.split():
         if (
             len(string_utils.strip_ansi_sequences(line))
@@ -235,11 +239,14 @@ def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str:
         ) > width:
             line = line[1:]
             line = justify_string(line, width=width, alignment=alignment)
-            retval = retval + "\n" + line
-            line = ""
-        line = line + " " + word
+            retval = retval + '\n' + line
+            line = indent
+        line = line + ' ' + word
     if len(string_utils.strip_ansi_sequences(line)) > 0:
-        retval += "\n" + line[1:]
+        if alignment != 'j':
+            retval += "\n" + justify_string(line[1:], width=width, alignment=alignment)
+        else:
+            retval += "\n" + line[1:]
     return retval[1:]