diff --git a/icecream/icecream.py b/icecream/icecream.py index 5a7ae4c..a87648d 100644 --- a/icecream/icecream.py +++ b/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 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 diff --git a/setup.py b/setup.py index df3d31d..38b6d02 100644 --- a/setup.py +++ b/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, diff --git a/tests/test_icecream.py b/tests/test_icecream.py index 8c695ee..e30d090 100644 --- a/tests/test_icecream.py +++ b/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())