gh-141510: Support frozendict in pprint (#144908)

Co-authored-by: devdanzin <74280297+devdanzin@users.noreply.github.com>

diff --git a/Lib/pprint.py b/Lib/pprint.py
index 92a2c543ac279cdbe92af1d7841c527929a500dc..e111bd59d4152c6c4a977ecd248fc283ca911a54 100644 (file)
--- a/Lib/pprint.py
+++ b/Lib/pprint.py
@@ -235,6 +235,20 @@ class PrettyPrinter:
 
     _dispatch[dict.__repr__] = _pprint_dict
 
+    def _pprint_frozendict(self, object, stream, indent, allowance, context, level):
+        write = stream.write
+        cls = object.__class__
+        stream.write(cls.__name__ + '(')
+        length = len(object)
+        if length:
+            self._pprint_dict(object, stream,
+                              indent + len(cls.__name__) + 1,
+                              allowance + 1,
+                              context, level)
+        write(')')
+
+    _dispatch[frozendict.__repr__] = _pprint_frozendict
+
     def _pprint_ordered_dict(self, object, stream, indent, allowance, context, level):
         if not len(object):
             stream.write(repr(object))
@@ -623,12 +637,21 @@ class PrettyPrinter:
             else:
                 return repr(object), True, False
 
-        if issubclass(typ, dict) and r is dict.__repr__:
+        if ((issubclass(typ, dict) and r is dict.__repr__)
+            or (issubclass(typ, frozendict) and r is frozendict.__repr__)):
+            is_frozendict = issubclass(typ, frozendict)
             if not object:
-                return "{}", True, False
+                if is_frozendict:
+                    rep = f"{object.__class__.__name__}()"
+                else:
+                    rep = "{}"
+                return rep, True, False
             objid = id(object)
             if maxlevels and level >= maxlevels:
-                return "{...}", False, objid in context
+                rep = "{...}"
+                if is_frozendict:
+                    rep = f"{object.__class__.__name__}({rep})"
+                return rep, False, objid in context
             if objid in context:
                 return _recursion(object), False, True
             context[objid] = 1
@@ -651,7 +674,10 @@ class PrettyPrinter:
                 if krecur or vrecur:
                     recursive = True
             del context[objid]
-            return "{%s}" % ", ".join(components), readable, recursive
+            rep = "{%s}" % ", ".join(components)
+            if is_frozendict:
+                rep = f"{object.__class__.__name__}({rep})"
+            return rep, readable, recursive
 
         if (issubclass(typ, list) and r is list.__repr__) or \
            (issubclass(typ, tuple) and r is tuple.__repr__):

diff --git a/Lib/test/test_pprint.py b/Lib/test/test_pprint.py
index 41c337ade7eca136c50706aa76a4fc0df8bdc62a..f3860a5d511989d7acb5b09566407167be0a8a12 100644 (file)
--- a/Lib/test/test_pprint.py
+++ b/Lib/test/test_pprint.py
@@ -67,6 +67,13 @@ class dict3(dict):
     def __repr__(self):
         return dict.__repr__(self)
 
+class frozendict2(frozendict):
+    pass
+
+class frozendict3(frozendict):
+    def __repr__(self):
+        return frozendict.__repr__(self)
+
 class dict_custom_repr(dict):
     def __repr__(self):
         return '*'*len(dict.__repr__(self))
@@ -254,18 +261,22 @@ class QueryTestCase(unittest.TestCase):
                        set(), set2(), set3(),
                        frozenset(), frozenset2(), frozenset3(),
                        {}, dict2(), dict3(),
+                       frozendict(), frozendict2(), frozendict3(),
                        {}.keys(), {}.values(), {}.items(),
                        MappingView({}), KeysView({}), ItemsView({}), ValuesView({}),
                        self.assertTrue, pprint,
                        -6, -6, -6-6j, -1.5, "x", b"x", bytearray(b"x"),
                        (3,), [3], {3: 6},
-                       (1,2), [3,4], {5: 6},
+                       (1,2), [3,4],
                        tuple2((1,2)), tuple3((1,2)), tuple3(range(100)),
                        [3,4], list2([3,4]), list3([3,4]), list3(range(100)),
                        set({7}), set2({7}), set3({7}),
                        frozenset({8}), frozenset2({8}), frozenset3({8}),
-                       dict2({5: 6}), dict3({5: 6}),
+                       {5: 6}, dict2({5: 6}), dict3({5: 6}),
+                       frozendict({5: 6}), frozendict2({5: 6}), frozendict3({5: 6}),
                        {5: 6}.keys(), {5: 6}.values(), {5: 6}.items(),
+                       frozendict({5: 6}).keys(), frozendict({5: 6}).values(),
+                       frozendict({5: 6}).items(),
                        MappingView({5: 6}), KeysView({5: 6}),
                        ItemsView({5: 6}), ValuesView({5: 6}),
                        range(10, -11, -1),
@@ -330,20 +341,45 @@ class QueryTestCase(unittest.TestCase):
         for type in [dict, dict2]:
             self.assertEqual(pprint.pformat(type(o)), exp)
 
+        exp = """\
+frozendict({'RPM_cal': 0,
+            'RPM_cal2': 48059,
+            'Speed_cal': 0,
+            'controldesk_runtime_us': 0,
+            'main_code_runtime_us': 0,
+            'read_io_runtime_us': 0,
+            'write_io_runtime_us': 43690})"""
+        self.assertEqual(pprint.pformat(frozendict(o)), exp)
+        exp = """\
+frozendict2({'RPM_cal': 0,
+             'RPM_cal2': 48059,
+             'Speed_cal': 0,
+             'controldesk_runtime_us': 0,
+             'main_code_runtime_us': 0,
+             'read_io_runtime_us': 0,
+             'write_io_runtime_us': 43690})"""
+        self.assertEqual(pprint.pformat(frozendict2(o)), exp)
+
         o = range(100)
         exp = 'dict_keys([%s])' % ',\n '.join(map(str, o))
         keys = dict.fromkeys(o).keys()
         self.assertEqual(pprint.pformat(keys), exp)
+        keys = frozendict.fromkeys(o).keys()
+        self.assertEqual(pprint.pformat(keys), exp)
 
         o = range(100)
         exp = 'dict_values([%s])' % ',\n '.join(map(str, o))
         values = {v: v for v in o}.values()
         self.assertEqual(pprint.pformat(values), exp)
+        values = frozendict({v: v for v in o}).values()
+        self.assertEqual(pprint.pformat(values), exp)
 
         o = range(100)
         exp = 'dict_items([%s])' % ',\n '.join("(%s, %s)" % (i, i) for i in o)
         items = {v: v for v in o}.items()
         self.assertEqual(pprint.pformat(items), exp)
+        items = frozendict({v: v for v in o}).items()
+        self.assertEqual(pprint.pformat(items), exp)
 
         o = range(100)
         exp = 'odict_keys([%s])' % ',\n '.join(map(str, o))