bpo-43080: pprint for dataclass instances (GH-24389)

* Added pprint support for dataclass instances which don't have a custom __repr__.

2 files changed, 123 insertions, 14 deletions

diff --git a/Lib/pprint.py b/Lib/pprint.py
index b45cfdd..13819f3 100644
--- a/Lib/pprint.py
+++ b/Lib/pprint.py
@@ -35,6 +35,7 @@ saferepr()
 """
 
 import collections as _collections
+import dataclasses as _dataclasses
 import re
 import sys as _sys
 import types as _types
@@ -178,8 +179,26 @@ class PrettyPrinter:
                 p(self, object, stream, indent, allowance, context, level + 1)
                 del context[objid]
                 return
+            elif (_dataclasses.is_dataclass(object) and
+                  not isinstance(object, type) and
+                  object.__dataclass_params__.repr and
+                  # Check dataclass has generated repr method.
+                  hasattr(object.__repr__, "__wrapped__") and
+                  "__create_fn__" in object.__repr__.__wrapped__.__qualname__):
+                context[objid] = 1
+                self._pprint_dataclass(object, stream, indent, allowance, context, level + 1)
+                del context[objid]
+                return
         stream.write(rep)
 
+    def _pprint_dataclass(self, object, stream, indent, allowance, context, level):
+        cls_name = object.__class__.__name__
+        indent += len(cls_name) + 1
+        items = [(f.name, getattr(object, f.name)) for f in _dataclasses.fields(object) if f.repr]
+        stream.write(cls_name + '(')
+        self._format_namespace_items(items, stream, indent, allowance, context, level)
+        stream.write(')')
+
     _dispatch = {}
 
     def _pprint_dict(self, object, stream, indent, allowance, context, level):
@@ -346,21 +365,9 @@ class PrettyPrinter:
         else:
             cls_name = object.__class__.__name__
         indent += len(cls_name) + 1
-        delimnl = ',\n' + ' ' * indent
         items = object.__dict__.items()
-        last_index = len(items) - 1
-
         stream.write(cls_name + '(')
-        for i, (key, ent) in enumerate(items):
-            stream.write(key)
-            stream.write('=')
-
-            last = i == last_index
-            self._format(ent, stream, indent + len(key) + 1,
-                         allowance if last else 1,
-                         context, level)
-            if not last:
-                stream.write(delimnl)
+        self._format_namespace_items(items, stream, indent, allowance, context, level)
         stream.write(')')
 
     _dispatch[_types.SimpleNamespace.__repr__] = _pprint_simplenamespace
@@ -382,6 +389,25 @@ class PrettyPrinter:
             if not last:
                 write(delimnl)
 
+    def _format_namespace_items(self, items, stream, indent, allowance, context, level):
+        write = stream.write
+        delimnl = ',\n' + ' ' * indent
+        last_index = len(items) - 1
+        for i, (key, ent) in enumerate(items):
+            last = i == last_index
+            write(key)
+            write('=')
+            if id(ent) in context:
+                # Special-case representation of recursion to match standard
+                # recursive dataclass repr.
+                write("...")
+            else:
+                self._format(ent, stream, indent + len(key) + 1,
+                             allowance if last else 1,
+                             context, level)
+            if not last:
+                write(delimnl)
+
     def _format_items(self, items, stream, indent, allowance, context, level):
         write = stream.write
         indent += self._indent_per_level
diff --git a/Lib/test/test_pprint.py b/Lib/test/test_pprint.py
index e5d2ac5..6c714fd 100644
--- a/Lib/test/test_pprint.py
+++ b/Lib/test/test_pprint.py
@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import collections
+import dataclasses
 import io
 import itertools
 import pprint
@@ -66,6 +67,38 @@ class dict_custom_repr(dict):
     def __repr__(self):
         return '*'*len(dict.__repr__(self))
 
+@dataclasses.dataclass
+class dataclass1:
+    field1: str
+    field2: int
+    field3: bool = False
+    field4: int = dataclasses.field(default=1, repr=False)
+
+@dataclasses.dataclass
+class dataclass2:
+    a: int = 1
+    def __repr__(self):
+        return "custom repr that doesn't fit within pprint width"
+
+@dataclasses.dataclass(repr=False)
+class dataclass3:
+    a: int = 1
+
+@dataclasses.dataclass
+class dataclass4:
+    a: "dataclass4"
+    b: int = 1
+
+@dataclasses.dataclass
+class dataclass5:
+    a: "dataclass6"
+    b: int = 1
+
+@dataclasses.dataclass
+class dataclass6:
+    c: "dataclass5"
+    d: int = 1
+
 class Unorderable:
     def __repr__(self):
         return str(id(self))
@@ -428,7 +461,7 @@ mappingproxy(OrderedDict([('the', 0),
             lazy=7,
             dog=8,
         )
-        formatted = pprint.pformat(ns, width=60)
+        formatted = pprint.pformat(ns, width=60, indent=4)
         self.assertEqual(formatted, """\
 namespace(the=0,
           quick=1,
@@ -465,6 +498,56 @@ AdvancedNamespace(the=0,
                   lazy=7,
                   dog=8)""")
 
+    def test_empty_dataclass(self):
+        dc = dataclasses.make_dataclass("MyDataclass", ())()
+        formatted = pprint.pformat(dc)
+        self.assertEqual(formatted, "MyDataclass()")
+
+    def test_small_dataclass(self):
+        dc = dataclass1("text", 123)
+        formatted = pprint.pformat(dc)
+        self.assertEqual(formatted, "dataclass1(field1='text', field2=123, field3=False)")
+
+    def test_larger_dataclass(self):
+        dc = dataclass1("some fairly long text", int(1e10), True)
+        formatted = pprint.pformat([dc, dc], width=60, indent=4)
+        self.assertEqual(formatted, """\
+[   dataclass1(field1='some fairly long text',
+               field2=10000000000,
+               field3=True),
+    dataclass1(field1='some fairly long text',
+               field2=10000000000,
+               field3=True)]""")
+
+    def test_dataclass_with_repr(self):
+        dc = dataclass2()
+        formatted = pprint.pformat(dc, width=20)
+        self.assertEqual(formatted, "custom repr that doesn't fit within pprint width")
+
+    def test_dataclass_no_repr(self):
+        dc = dataclass3()
+        formatted = pprint.pformat(dc, width=10)
+        self.assertRegex(formatted, r"<test.test_pprint.dataclass3 object at \w+>")
+
+    def test_recursive_dataclass(self):
+        dc = dataclass4(None)
+        dc.a = dc
+        formatted = pprint.pformat(dc, width=10)
+        self.assertEqual(formatted, """\
+dataclass4(a=...,
+           b=1)""")
+
+    def test_cyclic_dataclass(self):
+        dc5 = dataclass5(None)
+        dc6 = dataclass6(None)
+        dc5.a = dc6
+        dc6.c = dc5
+        formatted = pprint.pformat(dc5, width=10)
+        self.assertEqual(formatted, """\
+dataclass5(a=dataclass6(c=...,
+                        d=1),
+           b=1)""")
+
     def test_subclassing(self):
         # length(repr(obj)) > width
         o = {'names with spaces': 'should be presented using repr()',