Various changes.

[python_utils.git] / text_utils.py

diff --git a/text_utils.py b/text_utils.py
index 76b5db600ec588e10e6d9c951fc1bb9b3960d70e..3be32ff49ec05b2d7ca0978e6cb34b65da64162e 100644 (file)
--- a/text_utils.py
+++ b/text_utils.py
@@ -5,7 +5,7 @@
 from collections import defaultdict
 import math
 import sys
-from typing import List, NamedTuple
+from typing import List, NamedTuple, Optional
 
 from ansi import fg, reset
 
@@ -167,3 +167,52 @@ def generate_padded_columns(text: List[str]) -> str:
             word = justify_string(word, width=width, alignment='l')
             out += f'{word} '
         yield out
+
+
+def wrap_string(text: str, n: int) -> str:
+    chunks = text.split()
+    out = ''
+    width = 0
+    for chunk in chunks:
+        if width + len(chunk) > n:
+            out += '\n'
+            width = 0
+        out += chunk + ' '
+        width += len(chunk) + 1
+    return out
+
+
+class Indenter:
+    """
+    with Indenter(pad_count = 8) as i:
+        i.print('test')
+        with i:
+            i.print('-ing')
+            with i:
+                i.print('1, 2, 3')
+    """
+    def __init__(self,
+                 *,
+                 pad_prefix: Optional[str] = None,
+                 pad_char: str = ' ',
+                 pad_count: int = 4):
+        self.level = -1
+        if pad_prefix is not None:
+            self.pad_prefix = pad_prefix
+        else:
+            self.pad_prefix = ''
+        self.padding = pad_char * pad_count
+
+    def __enter__(self):
+        self.level += 1
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        self.level -= 1
+        if self.level < -1:
+            self.level = -1
+
+    def print(self, *arg, **kwargs):
+        import string_utils
+        text = string_utils.sprintf(*arg, **kwargs)
+        print(self.pad_prefix + self.padding * self.level + text, end='')