pretty printing output

Author: samuelcolvin

devtools/__init__.py

@@ -1,3 +1,5 @@
 # flake8: noqa
+from .ansi import *
 from .debug import *
+from .prettier import *
 from .version import VERSION

devtools/ansi.py

@@ -1,16 +1,26 @@
+import re
 import sys
 from enum import IntEnum
+from typing import Any
 
 _ansi_template = '\033[{}m'
+_ansi_re = re.compile('\033\[((?:\d|;)*)([a-zA-Z])')
 
+__all__ = ['sformat', 'sprint']
 
-def isatty(stream):
+
+def isatty(stream=None):
+    stream = stream or sys.stdout
     try:
         return stream.isatty()
     except Exception:
         return False
 
 
+def strip_ansi(value):
+    return _ansi_re.sub('', value)
+
+
 class Style(IntEnum):
     """
     Heavily borrowed from https://github.com/pallets/click/blob/6.7/click/termui.py
@@ -62,7 +72,7 @@ class Style(IntEnum):
     # this is a meta value used for the "Style" instance which is the "style" function
     function = -1
 
-    def __call__(self, text: str, *styles, reset: bool=True):
+    def __call__(self, text: Any, *styles, reset: bool=True):
         """
         Styles a text with ANSI styles and returns the new string.
 
@@ -103,9 +113,9 @@ def __call__(self, text: str, *styles, reset: bool=True):
             codes.append(str(s.value))
 
         if codes:
-            r = _ansi_template.format(';'.join(codes)) + text
+            r = _ansi_template.format(';'.join(codes)) + str(text)
         else:
-            r = text
+            r = str(text)
 
         if reset:
             r += _ansi_template.format(self.reset)
@@ -128,10 +138,10 @@ def __str__(self):
             return super().__str__()
 
 
-style = Style(-1)
+sformat = Style(-1)
 
 
 def sprint(text, *styles, reset=True, flush=True, file=None, **print_kwargs):
-    if isatty(file or sys.stdout):
-        text = style(text, *styles, reset=reset)
+    if isatty(file):
+        text = sformat(text, *styles, reset=reset)
     print(text, flush=flush, file=file, **print_kwargs)

devtools/debug.py

@@ -7,8 +7,12 @@
 from textwrap import dedent
 from typing import Generator, List, Optional, Tuple
 
+from .ansi import isatty, sformat
+from .prettier import PrettyFormat
+
 __all__ = ['Debug', 'debug']
 CWD = Path('.').resolve()
+pformat = PrettyFormat()
 
 
 def env_true(var_name, alt='FALSE'):
@@ -26,18 +30,13 @@ def __init__(self, value, *, name=None, **extra):
             self.extra.append(('len', len(value)))
         self.extra += [(k, v) for k, v in extra.items() if v is not None]
 
-    def value_str(self):
-        if isinstance(self.value, str):
-            return '"{}"'.format(self.value)
-        return str(self.value)
-
     def __str__(self) -> str:
         template = '{value} ({self.value.__class__.__name__}) {extra}'
         if self.name:
             template = '{self.name} = ' + template
         return template.format(
             self=self,
-            value=self.value_str(),
+            value=pformat(self.value, indent=2),
             extra=' '.join('{}={}'.format(k, v) for k, v in self.extra)
         ).rstrip(' ')  # trailing space if extra is empty
 
@@ -55,9 +54,19 @@ def __init__(self, *, filename: str, lineno: int, frame: str, arguments: List[De
         self.frame = frame
         self.arguments = arguments
 
+    def str(self, colours=False) -> str:
+        if colours:
+            prefix = '{}:{} {}\n  '.format(
+                sformat(self.filename, sformat.magenta),
+                sformat(self.lineno, sformat.green),
+                sformat(self.frame, sformat.green, sformat.italic)
+            )
+        else:
+            prefix = '{0.filename}:{0.lineno} {0.frame}\n  '.format(self)
+        return prefix + '\n  '.join(str(a) for a in self.arguments)
+
     def __str__(self) -> str:
-        template = '{s.filename}:{s.lineno} {s.frame}\n  {a}'
-        return template.format(s=self, a='\n  '.join(str(a) for a in self.arguments))
+        return self.str()
 
     def __repr__(self) -> str:
         arguments = ' '.join(str(a) for a in self.arguments)
@@ -88,13 +97,15 @@ def _env_bool(cls, value, env_name, env_default='TRUE'):
         else:
             return value
 
-    def __call__(self, *args, **kwargs):
-        print(self._process(args, kwargs, r'debug *\('), flush=True)
+    def __call__(self, *args, file_=None, flush_=True, **kwargs) -> None:
+        d_out = self._process(args, kwargs, r'debug *\(')
+        s = d_out.str(isatty(file_))
+        print(s, file=file_, flush=flush_)
 
-    def format(self, *args, **kwargs):
+    def format(self, *args, **kwargs) -> DebugOutput:
         return self._process(args, kwargs, r'debug.format *\(')
 
-    def _process(self, args, kwargs, func_regex):
+    def _process(self, args, kwargs, func_regex) -> DebugOutput:
         curframe = inspect.currentframe()
         frames = inspect.getouterframes(curframe, context=self._frame_context_length)
         # BEWARE: this must be call by a method which in turn is called "directly" for the frame to be correct

devtools/prettier.py

@@ -3,6 +3,8 @@
 import textwrap
 from typing import Any, Generator, Union
 
+from .ansi import isatty
+
 try:
     from pygments import highlight
     from pygments.lexers import PythonLexer
@@ -16,6 +18,7 @@
     (list, '[', ']'),
     (set, '{', '}'),
 ]
+__all__ = ['PrettyFormat', 'pformat', 'pprint']
 
 
 class PrettyFormat:
@@ -42,9 +45,9 @@ def __init__(self,
             (collections.Generator, self._format_generators),
         ]
 
-    def __call__(self, value: Any, *, indent_current: int=0):
+    def __call__(self, value: Any, *, indent: int=0, indent_first: bool=False):
         self._stream = io.StringIO()
-        self._format(value, indent_current=indent_current, indent_first=True)
+        self._format(value, indent_current=indent, indent_first=indent_first)
         s = self._stream.getvalue()
         if self._colorize and pyg_lexer:
             s = highlight(s, lexer=pyg_lexer, formatter=pyg_formatter)
@@ -142,7 +145,7 @@ def _format_raw(self, value: Any, value_repr: str, indent_current: int, indent_n
 
 
 pformat = PrettyFormat()
-_ppformat = PrettyFormat(colorize=True)
+_ppformat = PrettyFormat(colorize=isatty())
 
 
 def pprint(s):

tests/test_ansi.py

@@ -2,43 +2,48 @@
 
 import pytest
 
-from devtools.ansi import sprint, style
+from devtools.ansi import sformat, sprint
 
 
 def test_colorize():
-    v = style('hello', style.red)
+    v = sformat('hello', sformat.red)
     assert v == '\x1b[31mhello\x1b[0m'
 
 
 def test_no_reset():
-    v = style('hello', style.bold, reset=False)
+    v = sformat('hello', sformat.bold, reset=False)
     assert v == '\x1b[1mhello'
 
 
 def test_combine_styles():
-    v = style('hello', style.red, style.bold)
+    v = sformat('hello', sformat.red, sformat.bold)
     assert v == '\x1b[31;1mhello\x1b[0m'
 
 
 def test_no_styles():
-    v = style('hello')
+    v = sformat('hello')
     assert v == 'hello\x1b[0m'
 
 
 def test_style_str():
-    v = style('hello', 'red')
+    v = sformat('hello', 'red')
     assert v == '\x1b[31mhello\x1b[0m'
 
 
+def test_non_str_input():
+    v = sformat(12.2, sformat.yellow, sformat.italic)
+    assert v == '\x1b[33;3m12.2\x1b[0m'
+
+
 def test_invalid_style_str():
     with pytest.raises(ValueError) as exc_info:
-        style('x', 'mauve')
+        sformat('x', 'mauve')
     assert exc_info.value.args[0] == 'invalid style "mauve"'
 
 
 def test_print_not_tty():
     stream = io.StringIO()
-    sprint('hello', style.green, file=stream)
+    sprint('hello', sformat.green, file=stream)
     out = stream.getvalue()
     assert out == 'hello\n'
 
@@ -49,7 +54,7 @@ def isatty(self):
             return True
 
     stream = TTYStream()
-    sprint('hello', style.green, file=stream)
+    sprint('hello', sformat.green, file=stream)
     out = stream.getvalue()
     assert out == '\x1b[32mhello\x1b[0m\n', repr(out)
 
@@ -60,21 +65,21 @@ def isatty(self):
             raise RuntimeError('boom')
 
     stream = TTYStream()
-    sprint('hello', style.green, file=stream)
+    sprint('hello', sformat.green, file=stream)
     out = stream.getvalue()
     assert out == 'hello\n'
 
 
 def test_get_styles():
-    assert style.styles['bold'] == 1
-    assert style.styles['un_bold'] == 22
+    assert sformat.styles['bold'] == 1
+    assert sformat.styles['un_bold'] == 22
 
 
 def test_repr():
-    assert repr(style) == '<pseudo function style(text, *styles)>'
-    assert repr(style.red) == '<Style.red: 31>'
+    assert repr(sformat) == '<pseudo function style(text, *styles)>'
+    assert repr(sformat.red) == '<Style.red: 31>'
 
 
 def test_str():
-    assert str(style) == '<pseudo function style(text, *styles)>'
-    assert str(style.red) == 'Style.red'
+    assert str(sformat) == '<pseudo function style(text, *styles)>'
+    assert str(sformat.red) == 'Style.red'

tests/test_expr_render.py

@@ -40,17 +40,25 @@ def test_exotic_types():
     print(s)
     # list and generator comprehensions are wrong because ast is wrong, see https://bugs.python.org/issue31241
     assert (
-        'tests/test_expr_render.py:<line no> test_exotic_types\n'
-        '  sum(aa) = 6 (int)\n'
-        '  1 == 2 = False (bool)\n'
-        '  1 < 2 = True (bool)\n'
-        '  1 << 2 = 4 (int)\n'
-        '  \'t\' if True else \'f\' = "t" (str) len=1\n'
-        '  1 or 2 = 1 (int)\n'
-        '  [a for a in aa] = [1, 2, 3] (list)\n'
-        '  {a for a in aa} = {1, 2, 3} (set)\n'
-        '  {a: a + 1 for a in aa} = {1: 2, 2: 3, 3: 4} (dict)\n'
-        '  (a for a in aa) = <generator object test_exotic_types.<locals>.<genexpr> at 0x<hash>> (generator)'
+        "tests/test_expr_render.py:<line no> test_exotic_types\n"
+        "  sum(aa) = 6 (int)\n"
+        "  1 == 2 = False (bool)\n"
+        "  1 < 2 = True (bool)\n"
+        "  1 << 2 = 4 (int)\n"
+        "  't' if True else 'f' = 't' (str) len=1\n"
+        "  1 or 2 = 1 (int)\n"
+        "  [a for a in aa] = [1, 2, 3] (list)\n"
+        "  {a for a in aa} = {1, 2, 3} (set)\n"
+        "  {a: a + 1 for a in aa} = {\n"
+        "      1: 2,\n"
+        "      2: 3,\n"
+        "      3: 4,\n"
+        "  } (dict)\n"
+        "  (a for a in aa) = (\n"
+        "      1,\n"
+        "      2,\n"
+        "      3,\n"
+        "  ) (generator)"
     ) == s