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,10 @@ class Conclusion:
     status: Status = attr.ib()
     interesting_origin: Optional[InterestingOrigin] = attr.ib()
 
+    def _repr_pretty_(self, p, cycle):
+        assert cycle is False
+        p.text(f"Conclusion ({self.status!r})")
+
 
 # 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 +516,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 +935,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/cover/test_pretty.py

@@ -56,6 +56,7 @@
 import pytest
 
 from hypothesis.internal.compat import PYPY
+from hypothesis.internal.conjecture.datatree import DataTree
 from hypothesis.internal.floats import SIGNALING_NAN
 from hypothesis.vendor import pretty
 
@@ -692,3 +693,10 @@ def test_pretty_partial_with_cycle():
     assert pretty.pretty(p) == "functools.partial(bool, [])"
     ls.append(p)
     assert pretty.pretty(p) == "functools.partial(bool, [functools.partial(bool, ...)])"
+
+
+def test_datatree_repr():
+    tree = DataTree()
+    observer = tree.new_observer()
+    observer.draw_boolean(True, was_forced=False, kwargs={"p": 0.5})
+    assert pretty.pretty(tree) == "boolean True {'p': 0.5}"