Initial version of auto-detection of terminal width.

icecream/icecream.py

@@ -15,6 +15,7 @@
 
 import ast
 import inspect
+import os
 import pprint
 import sys
 from datetime import datetime
@@ -33,6 +34,15 @@
 
 from .coloring import SolarizedDark
 
+try:
+    from shutil import get_terminal_size
+except ImportError:
+    try:
+        from backports.shutil_get_terminal_size import get_terminal_size
+    except ImportError:
+        def get_terminal_size():
+            return os.environ['COLUMNS']
+
 
 PYTHON2 = (sys.version_info[0] == 2)
 
@@ -83,6 +93,7 @@ def colorizedStderrPrint(s):
 
 
 DEFAULT_PREFIX = 'ic| '
+DEFAULT_TERMINAL_WIDTH = 80
 DEFAULT_LINE_WRAP_WIDTH = 70  # Characters.
 DEFAULT_CONTEXT_DELIMITER = '- '
 DEFAULT_OUTPUT_FUNCTION = colorizedStderrPrint
@@ -154,25 +165,51 @@ def format_pair(prefix, arg, value):
     return '\n'.join(lines)
 
 
-def argumentToString(obj):
-    s = DEFAULT_ARG_TO_STRING_FUNCTION(obj)
+def argumentToString(obj, width=DEFAULT_LINE_WRAP_WIDTH):
+    s = DEFAULT_ARG_TO_STRING_FUNCTION(obj, width=width)
     s = s.replace('\\n', '\n')  # Preserve string newlines in output.
     return s
 
 
+def detect_terminal_width(default=DEFAULT_TERMINAL_WIDTH):
+    """ Returns the number of columns that this terminal can handle. """
+    try:
+        # We need to pass a terminal height in the tuple so we pass the default
+        # of 25 lines but it's not used for anything.
+        width = get_terminal_size((default, 25)).columns
+    except Exception:  # Not in TTY or something else went wrong
+        width = default
+    # TODO account for argPrefix()
+    return width
+
+
+def supports_param(fn, param="width"):
+    """ Returns True if the function supports that parameter. """
+    try:
+        from inspect import signature
+        return param in signature(fn).parameters
+    except ImportError:  # Python 2.x
+        from inspect import getargspec
+        return param in getargspec(fn).args
+
+
 class IceCreamDebugger:
     _pairDelimiter = ', '  # Used by the tests in tests/.
-    lineWrapWidth = DEFAULT_LINE_WRAP_WIDTH
     contextDelimiter = DEFAULT_CONTEXT_DELIMITER
+    terminalWidth = DEFAULT_TERMINAL_WIDTH
+    lineWrapWidth = DEFAULT_LINE_WRAP_WIDTH
 
     def __init__(self, prefix=DEFAULT_PREFIX,
                  outputFunction=DEFAULT_OUTPUT_FUNCTION,
-                 argToStringFunction=argumentToString, includeContext=False):
+                 argToStringFunction=argumentToString, includeContext=False,
+                 detectTerminalWidth=False):
         self.enabled = True
         self.prefix = prefix
         self.includeContext = includeContext
         self.outputFunction = outputFunction
         self.argToStringFunction = argToStringFunction
+        self.passWidthParam = supports_param(self.argToStringFunction)
+        self._setLineWrapWidth(detectTerminalWidth=detectTerminalWidth)
 
     def __call__(self, *args):
         if self.enabled:
@@ -193,6 +230,17 @@ def __call__(self, *args):
 
         return passthrough
 
+    def _setLineWrapWidth(self, detectTerminalWidth=False, terminalWidth=None):
+        prefix_length = len(self.prefix()) if callable(self.prefix) else len(self.prefix)
+        if terminalWidth:
+            width = terminalWidth
+        elif detectTerminalWidth is True:
+            width = detect_terminal_width(DEFAULT_TERMINAL_WIDTH)
+        else:
+            width = DEFAULT_TERMINAL_WIDTH
+        self.terminalWidth = width
+        self.lineWrapWidth = width - prefix_length
+
     def format(self, *args):
         callFrame = inspect.currentframe().f_back
         out = self._format(callFrame, *args)
@@ -232,7 +280,8 @@ def _constructArgumentOutput(self, prefix, context, pairs):
         def argPrefix(arg):
             return '%s: ' % arg
 
-        pairs = [(arg, self.argToStringFunction(val)) for arg, val in pairs]
+        kwargs = {"width": self.lineWrapWidth} if self.passWidthParam else {}
+        pairs = [(arg, self.argToStringFunction(val, **kwargs)) for arg, val in pairs]
         # For cleaner output, if <arg> is a literal, eg 3, "string", b'bytes',
         # etc, only output the value, not the argument and the value, as the
         # argument and the value will be identical or nigh identical. Ex: with
@@ -316,15 +365,21 @@ def disable(self):
         self.enabled = False
 
     def configureOutput(self, prefix=_absent, outputFunction=_absent,
-                        argToStringFunction=_absent, includeContext=_absent):
+                        argToStringFunction=_absent, includeContext=_absent,
+                        terminalWidth=_absent):
         if prefix is not _absent:
             self.prefix = prefix
 
+        if prefix is not _absent or terminalWidth is not _absent:
+            new_terminal_width = terminalWidth if terminalWidth is not _absent else None
+            self._setLineWrapWidth(new_terminal_width)
+
         if outputFunction is not _absent:
             self.outputFunction = outputFunction
 
         if argToStringFunction is not _absent:
             self.argToStringFunction = argToStringFunction
+            self.passWidthParam = supports_param(self.argToStringFunction)
 
         if includeContext is not _absent:
             self.includeContext = includeContext

setup.py

@@ -103,6 +103,7 @@ def run_tests(self):
         'pygments>=2.2.0',
         'executing>=0.3.1',
         'asttokens>=2.0.1',
+        'backports.shutil-get-terminal-size==1.0.0; python_version < "3.3.0"',
     ],
     cmdclass={
         'test': RunTests,

tests/test_icecream.py

@@ -9,6 +9,8 @@
 #
 # License: MIT
 #
+import os
+import textwrap
 
 import sys
 import unittest
@@ -22,11 +24,9 @@
 import icecream
 from icecream import ic, stderrPrint, NoSourceAvailableError
 
-
 TEST_PAIR_DELIMITER = '| '
 MYFILENAME = basename(__file__)
 
-
 a = 1
 b = 2
 c = 3
@@ -61,11 +61,13 @@ def disableColoring():
 
 @contextmanager
 def configureIcecreamOutput(prefix=None, outputFunction=None,
-                            argToStringFunction=None, includeContext=None):
+                            argToStringFunction=None, includeContext=None, terminalWidth=None):
+
     oldPrefix = ic.prefix
     oldOutputFunction = ic.outputFunction
     oldArgToStringFunction = ic.argToStringFunction
     oldIncludeContext = ic.includeContext
+    oldTerminalWidth = ic.terminalWidth
 
     if prefix:
         ic.configureOutput(prefix=prefix)
@@ -75,12 +77,26 @@ def configureIcecreamOutput(prefix=None, outputFunction=None,
         ic.configureOutput(argToStringFunction=argToStringFunction)
     if includeContext:
         ic.configureOutput(includeContext=includeContext)
+    if terminalWidth:
+        ic.configureOutput(terminalWidth=terminalWidth)
 
     yield
 
     ic.configureOutput(
         oldPrefix, oldOutputFunction, oldArgToStringFunction,
-        oldIncludeContext)
+        oldIncludeContext, oldTerminalWidth)
+
+
+@contextmanager
+def detectTerminalWidth(terminal_width=icecream.DEFAULT_TERMINAL_WIDTH):
+    width = str(terminal_width)
+    old_terminal_width = os.getenv('COLUMNS', width)
+    try:
+        os.environ['COLUMNS'] = width
+        yield ic._setLineWrapWidth(detectTerminalWidth=True)
+    finally:
+        os.environ['COLUMNS'] = old_terminal_width
+        ic._setLineWrapWidth(detectTerminalWidth=False, terminalWidth=int(old_terminal_width))
 
 
 @contextmanager
@@ -182,6 +198,8 @@ def parseOutputIntoPairs(out, err, assertNumLines,
 class TestIceCream(unittest.TestCase):
     def setUp(self):
         ic._pairDelimiter = TEST_PAIR_DELIMITER
+        ic.configureOutput(prefix=icecream.DEFAULT_PREFIX,
+                           terminalWidth=icecream.DEFAULT_TERMINAL_WIDTH)
 
     def testWithoutArgs(self):
         with disableColoring(), captureStandardStreams() as (out, err):
@@ -518,3 +536,78 @@ def testColoring(self):
             ic({1: 'str'})  # Output should be colored with ANSI control codes.
 
         assert hasAnsiEscapeCodes(err.getvalue())
+
+    def testStringWithShortLineWrapWidth(self):
+        """ Test a string with a short line wrap width. """
+        ic._setLineWrapWidth(terminalWidth=10)
+        s = "123456789 1234567890"
+        with disableColoring(), captureStandardStreams() as (out, err):
+            ic(s)
+        if icecream.PYTHON2:
+            expected = "ic| s: '123456789 1234567890'"
+        else:
+            expected = textwrap.dedent("""
+            ic| s: ('123456789 '
+                    '1234567890')
+            """).strip()
+        self.assertEqual(err.getvalue().strip(), expected)
+
+    def testListWithShortLineWrapWidth(self):
+        """ Test a list with a short line wrap width. """
+        ic._setLineWrapWidth(terminalWidth=10)
+        lst = ["1 2 3 4 5", "2", "3", "4"]
+        with disableColoring(), captureStandardStreams() as (out, err):
+            ic(lst)
+        if icecream.PYTHON2:
+            expected = textwrap.dedent("""
+                ic| lst: ['1 2 3 4 5',
+                          '2',
+                          '3',
+                          '4']""").strip()
+        else:
+            expected = textwrap.dedent("""
+                ic| lst: ['1 '
+                          '2 '
+                          '3 '
+                          '4 '
+                          '5',
+                          '2',
+                          '3',
+                          '4']""").strip()
+        self.assertEqual(err.getvalue().strip(), expected)
+
+    def testLiteralWithShortTerminalWidth(self):
+        """ Test a literal with a short line wrap width. """
+        with detectTerminalWidth(10):
+            with disableColoring(), captureStandardStreams() as (out, err):
+                ic("banana banana")
+        if icecream.PYTHON2:
+            expected = 'ic| "banana banana": \'banana banana\''
+        else:
+            expected = textwrap.dedent("""
+        ic| "banana banana": ('banana '
+                              'banana')""").strip()
+        actual = err.getvalue().strip()
+        self.assertEqual(expected, actual)
+
+    def testConfigureOutput(self):
+        """ Test that line width is adjusted after running configureOutput()
+            with a new prefix. ic.lineWrapWidth will start at 70 then adjust
+            to 60, so a string that didn't wrap before should wrap now.
+        """
+        s = "a 70 character string a 70 character string a 70 character string a 70"
+        with disableColoring(), captureStandardStreams() as (out, err):
+            ic(s)
+        self.assertEqual("ic| s: '%s'" % s, err.getvalue().strip())
+        with configureIcecreamOutput(prefix="10prefix| ",
+                                     outputFunction=stderrPrint,
+                                     terminalWidth=icecream.DEFAULT_TERMINAL_WIDTH):
+            with disableColoring(), captureStandardStreams() as (out, err):
+                ic(s)
+        if icecream.PYTHON2:
+            expected = "10prefix| s: '%s'" % s
+        else:
+            expected = textwrap.dedent("""
+            10prefix| s: ('a 70 character string a 70 character string a 70 character string '
+                          'a 70')""").strip()
+        self.assertEqual(expected, err.getvalue().strip())