add pretty printers for DataTree

hypothesis-python/src/hypothesis/internal/conjecture/datatree.py

@@ -63,6 +63,15 @@ class Killed:
 
     next_node = attr.ib()
 
+    def _repr_pretty_(self, p, cycle):
+        assert cycle is False
+        p.text("Killed")
+
+
+def _node_pretty(ir_type, value, kwargs, *, forced):
+    forced_marker = " [forced]" if forced else ""
+    return f"{ir_type} {value}{forced_marker} {kwargs}"
+
 
 @attr.s(slots=True)
 class Branch:
@@ -79,6 +88,16 @@ def max_children(self):
         assert max_children > 0
         return max_children
 
+    def _repr_pretty_(self, p, cycle):
+        assert cycle is False
+        for i, (value, child) in enumerate(self.children.items()):
+            if i > 0:
+                p.break_()
+            p.text(_node_pretty(self.ir_type, value, self.kwargs, forced=False))
+            with p.indent(2):
+                p.break_()
+                p.pretty(child)
+
 
 @attr.s(slots=True, frozen=True)
 class Conclusion:
@@ -87,6 +106,15 @@ class Conclusion:
     status: Status = attr.ib()
     interesting_origin: Optional[InterestingOrigin] = attr.ib()
 
+    def _repr_pretty_(self, p, cycle):
+        assert cycle is False
+        o = self.interesting_origin
+        # avoid str(o), which can include multiple lines of context
+        origin = (
+            "" if o is None else f", {o.exc_type.__name__} at {o.filename}:{o.lineno}"
+        )
+        p.text(f"Conclusion ({self.status!r}{origin})")
+
 
 # The number of max children where, beyond this, it is practically impossible
 # for hypothesis to saturate / explore all children nodes in a reasonable time
@@ -493,6 +521,29 @@ def check_exhausted(self):
                 )
         return self.is_exhausted
 
+    def _repr_pretty_(self, p, cycle):
+        assert cycle is False
+        indent = 0
+        for i, (ir_type, kwargs, value) in enumerate(
+            zip(self.ir_types, self.kwargs, self.values)
+        ):
+            with p.indent(indent):
+                if i > 0:
+                    p.break_()
+                p.text(_node_pretty(ir_type, value, kwargs, forced=i in self.forced))
+            indent += 2
+
+        if isinstance(self.transition, Branch):
+            if len(self.values) > 0:
+                p.break_()
+            p.pretty(self.transition)
+
+        if isinstance(self.transition, (Killed, Conclusion)):
+            with p.indent(indent):
+                if len(self.values) > 0:
+                    p.break_()
+                p.pretty(self.transition)
+
 
 class DataTree:
     """
@@ -889,6 +940,10 @@ def _reject_child(self, ir_type, kwargs, *, child, key):
         if child in children:
             children.remove(child)
 
+    def _repr_pretty_(self, p, cycle):
+        assert cycle is False
+        return p.pretty(self.root)
+
 
 class TreeRecordingObserver(DataObserver):
     def __init__(self, tree):

hypothesis-python/tests/conjecture/test_data_tree.py

@@ -8,6 +8,7 @@
 # v. 2.0. If a copy of the MPL was not distributed with this file, You can
 # obtain one at https://mozilla.org/MPL/2.0/.
 
+import textwrap
 from random import Random
 
 import pytest
@@ -22,7 +23,9 @@
 )
 from hypothesis.internal.conjecture.engine import ConjectureRunner
 from hypothesis.internal.conjecture.floats import float_to_int
+from hypothesis.internal.escalation import InterestingOrigin
 from hypothesis.internal.floats import next_up
+from hypothesis.vendor import pretty
 
 from tests.conjecture.common import (
     draw_boolean_kwargs,
@@ -567,3 +570,42 @@ def buf(data):
         prefix = tree.generate_novel_prefix(Random())
         data = ConjectureData.for_buffer(prefix)
         assert data.draw_float(min_value, max_value, allow_nan=False) == expected_value
+
+
+@given(draw_boolean_kwargs(), draw_integer_kwargs())
+def test_datatree_repr(bool_kwargs, int_kwargs):
+    tree = DataTree()
+
+    try:
+        int("not an int")
+    except ValueError as e:
+        origin = InterestingOrigin.from_exception(e)
+
+    observer = tree.new_observer()
+    observer.draw_boolean(True, was_forced=False, kwargs=bool_kwargs)
+    observer.conclude_test(Status.INVALID, interesting_origin=None)
+
+    observer = tree.new_observer()
+    observer.draw_boolean(False, was_forced=False, kwargs=bool_kwargs)
+    observer.draw_integer(42, was_forced=False, kwargs=int_kwargs)
+    observer.conclude_test(Status.VALID, interesting_origin=None)
+
+    observer = tree.new_observer()
+    observer.draw_boolean(False, was_forced=False, kwargs=bool_kwargs)
+    observer.draw_integer(0, was_forced=False, kwargs=int_kwargs)
+    observer.conclude_test(Status.INTERESTING, interesting_origin=origin)
+
+    assert (
+        pretty.pretty(tree)
+        == textwrap.dedent(
+            f"""
+        boolean True {bool_kwargs}
+          Conclusion (Status.INVALID)
+        boolean False {bool_kwargs}
+          integer 42 {int_kwargs}
+            Conclusion (Status.VALID)
+          integer 0 {int_kwargs}
+            Conclusion (Status.INTERESTING, {origin})
+        """
+        ).strip()
+    )