implement support for __pretty__ (#36)

* implment support for __pretty__

* fix tests, use dict

* uprev

* fix intermittent timing test

Author: samuelcolvin

HISTORY.rst

@@ -3,7 +3,12 @@
 History
 -------
 
-v0.4.0 (2017-12-29)
+v0.5.0 (2019-01-03)
+...................
+* support ``MultiDict``, #34
+* support ``__pretty__`` method, #36
+
+v0.4.0 (2018-12-29)
 ...................
 * remove use of ``warnings``, include in output, #30
 * fix rendering errors #31

devtools/prettier.py

@@ -4,6 +4,7 @@
 from collections import OrderedDict
 from collections.abc import Generator
 from typing import Any, Union
+from unittest.mock import _Call as MockCall
 
 from .ansi import isatty
 
@@ -29,6 +30,8 @@
     (frozenset, 'frozenset({', '})'),
 ]
 DEFAULT_WIDTH = int(os.getenv('PY_DEVTOOLS_WIDTH', 120))
+MISSING = object()
+PRETTY_KEY = '__prettier_formatted_value__'
 
 
 def env_true(var_name, alt=None):
@@ -39,6 +42,14 @@ def env_true(var_name, alt=None):
         return alt
 
 
+def fmt(v):
+    return {PRETTY_KEY: v}
+
+
+class SkipPretty(Exception):
+    pass
+
+
 class PrettyFormat:
     def __init__(self,
                  indent_step=4,
@@ -77,6 +88,16 @@ def _format(self, value: Any, indent_current: int, indent_first: bool):
         if indent_first:
             self._stream.write(indent_current * self._c)
 
+        pretty_func = getattr(value, '__pretty__', None)
+        if pretty_func and not isinstance(value, MockCall):
+            try:
+                gen = pretty_func(fmt=fmt, skip_exc=SkipPretty)
+                self._render_pretty(gen, indent_current)
+            except SkipPretty:
+                pass
+            else:
+                return
+
         value_repr = repr(value)
         if len(value_repr) <= self._simple_cutoff and not isinstance(value, Generator):
             self._stream.write(value_repr)
@@ -88,6 +109,26 @@ def _format(self, value: Any, indent_current: int, indent_first: bool):
                     return
             self._format_raw(value, value_repr, indent_current, indent_new)
 
+    def _render_pretty(self, gen, indent: int):
+        prefix = False
+        for v in gen:
+            if isinstance(v, int) and v in {-1, 0, 1}:
+                indent += v * self._indent_step
+                prefix = True
+            else:
+                if prefix:
+                    self._stream.write('\n' + self._c * indent)
+                    prefix = False
+
+                pretty_value = v.get(PRETTY_KEY, MISSING) if (isinstance(v, dict) and len(v) == 1) else MISSING
+                if pretty_value is not MISSING:
+                    self._format(pretty_value, indent, False)
+                elif isinstance(v, str):
+                    self._stream.write(v)
+                else:
+                    # shouldn't happen but will
+                    self._stream.write(repr(v))
+
     def _format_dict(self, value: dict, value_repr: str, indent_current: int, indent_new: int):
         open_, before_, split_, after_, close_ = '{\n', indent_new * self._c, ': ', ',\n', '}'
         if isinstance(value, OrderedDict):

devtools/version.py

@@ -2,4 +2,4 @@
 
 __all__ = ['VERSION']
 
-VERSION = StrictVersion('0.4.0')
+VERSION = StrictVersion('0.5.0')

tests/test_custom_pretty.py

@@ -0,0 +1,49 @@
+from devtools import pformat
+
+
+def test_simple():
+    class CustomCls:
+        def __pretty__(self, fmt, **kwargs):
+            yield 'Thing('
+            yield 1
+            for i in range(3):
+                yield fmt(list(range(i)))
+                yield ','
+                yield 0
+            yield -1
+            yield ')'
+
+    my_cls = CustomCls()
+
+    v = pformat(my_cls)
+    assert v == """\
+Thing(
+    [],
+    [0],
+    [0, 1],
+)"""
+
+
+def test_skip():
+    class CustomCls:
+        def __pretty__(self, fmt, skip_exc, **kwargs):
+            raise skip_exc()
+            yield 'Thing()'
+
+        def __repr__(self):
+            return '<CustomCls repr>'
+
+    my_cls = CustomCls()
+    v = pformat(my_cls)
+    assert v == '<CustomCls repr>'
+
+
+def test_yield_other():
+    class CustomCls:
+        def __pretty__(self, fmt, **kwargs):
+            yield fmt('xxx')
+            yield 123
+
+    my_cls = CustomCls()
+    v = pformat(my_cls)
+    assert v == "'xxx'123"

tests/test_timer.py

@@ -44,7 +44,7 @@ def test_multiple_not_verbose():
         with t(verbose=False):
             sleep(i)
     t.summary(True)
-    v = f.getvalue()
+    v = re.sub('[123]s', '0s', f.getvalue())
     assert v == (
         '    0.010s elapsed\n'
         '    0.020s elapsed\n'