From 34ec8484795ab497ccb734ae32e36715e475d674 Mon Sep 17 00:00:00 2001
From: Kentaro Wada <www.kentaro.wada@gmail.com>
Date: Fri, 11 Sep 2015 04:26:59 +0900
Subject: [PATCH] Support clint.textui.max_width as context manager

---
 clint/textui/core.py       |  7 ++++-
 clint/textui/formatters.py | 60 +++++++++++++++++++++++++++++++++-----
 2 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/clint/textui/core.py b/clint/textui/core.py
index b61e2c1..e3c4a46 100644
--- a/clint/textui/core.py
+++ b/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)
diff --git a/clint/textui/formatters.py b/clint/textui/formatters.py
index 9962351..549808b 100644
--- a/clint/textui/formatters.py
+++ b/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,7 +35,57 @@ 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
@@ -41,12 +93,6 @@ def max_width(string, cols, separator='\n'):
     :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)