Support clint.textui.max_width as context manager

clint/textui/core.py

@@ -15,7 +15,7 @@
 
 from contextlib import contextmanager
 
-from .formatters import max_width, min_width
+from .formatters import max_width, min_width, _get_max_width_context
 from .cols import columns
 from ..utils import tsplit
 
@@ -53,6 +53,11 @@ def _indent(indent=0, quote='', indent_char=' '):
 
 def puts(s='', newline=True, stream=STDOUT):
     """Prints given string to stdout."""
+    max_width_ctx = _get_max_width_context()
+    if max_width_ctx:
+        cols, separator = max_width_ctx[-1]
+        s = max_width(s, cols, separator)
+
     if newline:
         s = tsplit(s, NEWLINES)
         s = map(str, s)

clint/textui/formatters.py

@@ -9,12 +9,14 @@
 """
 
 from __future__ import absolute_import
+from contextlib import contextmanager
 
 from .colored import ColoredString, clean
 from ..utils import tsplit, schunk
 
 
 NEWLINES = ('\n', '\r', '\r\n')
+MAX_WIDTHS = []
 
 
 def min_width(string, cols, padding=' '):
@@ -33,20 +35,64 @@ def min_width(string, cols, padding=' '):
     return '\n'.join(stack)
 
 
-def max_width(string, cols, separator='\n'):
+def _get_max_width_context():
+    return MAX_WIDTHS
+
+@contextmanager
+def _max_width_context():
+    """Max width context manager."""
+    try:
+        yield
+    finally:
+        MAX_WIDTHS.pop()
+
+def max_width(*args, **kwargs):
+    """Returns formatted text or context manager for textui:puts.
+
+        >>> from clint.textui import puts, max_width
+        >>> max_width('123 5678', 8)
+        '123 5678'
+        >>> max_width('123 5678', 7)
+        '123 \n5678'
+        >>> with max_width(7):
+        ...     puts('123 5678')
+        '123 \n5678'
+    """
+    args = list(args)
+
+    if not args:
+        args.append(kwargs.get('string'))
+        args.append(kwargs.get('cols'))
+        args.append(kwargs.get('separator'))
+    elif len(args) == 1:
+        args.append(kwargs.get('cols'))
+        args.append(kwargs.get('separator'))
+    elif len(args) == 2:
+        args.append(kwargs.get('separator'))
+
+    string, cols, separator = args
+    if separator is None:
+        separator = '\n'  # default value
+    if cols is None:
+        # cols should be specified vitally
+        # because string can be specified at textui:puts function
+        string, cols = cols, string
+
+    if string is None:
+        MAX_WIDTHS.append((cols, separator))
+        return _max_width_context()
+    else:
+        return _max_width_formatter(string, cols, separator)
+
+
+def _max_width_formatter(string, cols, separator='\n'):
     """Returns a freshly formatted
     :param string: string to be formatted
     :type string: basestring or clint.textui.colorred.ColoredString
     :param cols: max width the text to be formatted
     :type cols: int
     :param separator: separator to break rows
     :type separator: basestring
-
-        >>> formatters.max_width('123 5678', 8)
-        '123 5678'
-        >>> formatters.max_width('123 5678', 7)
-        '123 \n5678'
-
     """
 
     is_color = isinstance(string, ColoredString)