Issue 3976: fix pprint for sets, frozensets, and dicts containing unorderable types.

2 files changed, 49 insertions, 11 deletions

diff --git a/Lib/pprint.py b/Lib/pprint.py
index 8bbc845..b3a6446 100644
--- a/Lib/pprint.py
+++ b/Lib/pprint.py
@@ -70,6 +70,32 @@ def isrecursive(object):
     """Determine if object requires a recursive representation."""
     return _safe_repr(object, {}, None, 0)[2]
 
+class _safe_key:
+    """Helper function for key functions when sorting unorderable objects.
+
+    The wrapped-object will fallback to an Py2.x style comparison for
+    unorderable types (sorting first comparing the type name and then by
+    the obj ids).  Does not work recursively, so dict.items() must have
+    _safe_key applied to both the key and the value.
+
+    """
+
+    __slots__ = ['obj']
+
+    def __init__(self, obj):
+        self.obj = obj
+
+    def __lt__(self, other):
+        rv = self.obj.__lt__(other.obj)
+        if rv is NotImplemented:
+            rv = (str(type(self.obj)), id(self.obj)) < \
+                 (str(type(other.obj)), id(other.obj))
+        return rv
+
+def _safe_tuple(t):
+    "Helper function for comparing 2-tuples"
+    return _safe_key(t[0]), _safe_key(t[1])
+
 class PrettyPrinter:
     def __init__(self, indent=1, width=80, depth=None, stream=None):
         """Handle pretty printing operations onto a stream using a set of
@@ -145,7 +171,7 @@ class PrettyPrinter:
                 if length:
                     context[objid] = 1
                     indent = indent + self._indent_per_level
-                    items  = sorted(object.items())
+                    items  = sorted(object.items(), key=_safe_tuple)
                     key, ent = items[0]
                     rep = self._repr(key, context, level)
                     write(rep)
@@ -178,14 +204,14 @@ class PrettyPrinter:
                         return
                     write('{')
                     endchar = '}'
-                    object = sorted(object)
+                    object = sorted(object, key=_safe_key)
                 elif issubclass(typ, frozenset):
                     if not length:
                         write('frozenset()')
                         return
                     write('frozenset({')
                     endchar = '})'
-                    object = sorted(object)
+                    object = sorted(object, key=_safe_key)
                     indent += 10
                 else:
                     write('(')
@@ -267,14 +293,7 @@ def _safe_repr(object, context, maxlevels, level):
         append = components.append
         level += 1
         saferepr = _safe_repr
-        items = object.items()
-        try:
-            items = sorted(items)
-        except TypeError:
-            def sortkey(item):
-                key, value = item
-                return str(type(key)), key, value
-            items = sorted(items, key=sortkey)
+        items = sorted(object.items(), key=_safe_tuple)
         for k, v in items:
             krepr, kreadable, krecur = saferepr(k, context, maxlevels, level)
             vrepr, vreadable, vrecur = saferepr(v, context, maxlevels, level)
diff --git a/Lib/test/test_pprint.py b/Lib/test/test_pprint.py
index 3d13b60..4370935 100644
--- a/Lib/test/test_pprint.py
+++ b/Lib/test/test_pprint.py
@@ -2,6 +2,7 @@ import pprint
 import test.support
 import unittest
 import test.test_set
+import random
 
 # list, tuple and dict subclasses that do or don't overwrite __repr__
 class list2(list):
@@ -25,6 +26,10 @@ class dict3(dict):
     def __repr__(self):
         return dict.__repr__(self)
 
+class Unorderable:
+    def __repr__(self):
+        return str(id(self))
+
 class QueryTestCase(unittest.TestCase):
 
     def setUp(self):
@@ -407,6 +412,20 @@ class QueryTestCase(unittest.TestCase):
         self.assertEqual(pprint.pformat(nested_dict, depth=1), lv1_dict)
         self.assertEqual(pprint.pformat(nested_list, depth=1), lv1_list)
 
+    def test_sort_unorderable_values(self):
+        # Issue 3976:  sorted pprints fail for unorderable values.
+        n = 20
+        keys = [Unorderable() for i in range(n)]
+        random.shuffle(keys)
+        skeys = sorted(keys, key=id)
+        clean = lambda s: s.replace(' ', '').replace('\n','')
+
+        self.assertEqual(clean(pprint.pformat(set(keys))),
+            '{' + ','.join(map(repr, skeys)) + '}')
+        self.assertEqual(clean(pprint.pformat(frozenset(keys))),
+            'frozenset({' + ','.join(map(repr, skeys)) + '})')
+        self.assertEqual(clean(pprint.pformat(dict.fromkeys(keys))),
+            '{' + ','.join('%r:None' % k for k in skeys) + '}')
 
 class DottedPrettyPrinter(pprint.PrettyPrinter):