pprint's workhorse _safe_repr() function took time quadratic in the # of

elements when crunching a list, dict or tuple.  Now takes linear time
instead -- huge speedup for even moderately large containers, and the
code is notably simpler too.
Added some basic "is the output correct?" tests to test_pprint.

2 files changed, 44 insertions, 39 deletions

diff --git a/Lib/pprint.py b/Lib/pprint.py
index 48efb33..351323b 100644
--- a/Lib/pprint.py
+++ b/Lib/pprint.py
@@ -188,66 +188,55 @@ class PrettyPrinter:
 # Return triple (repr_string, isreadable, isrecursive).
 
 def _safe_repr(object, context, maxlevels=None, level=0):
-    level = level + 1
+    level += 1
     typ = type(object)
     if not (typ in (DictType, ListType, TupleType) and object):
         rep = `object`
         return rep, (rep and (rep[0] != '<')), 0
+
     if context.has_key(id(object)):
         return `_Recursion(object)`, 0, 1
     objid = id(object)
     context[objid] = 1
+
     readable = 1
     recursive = 0
-    if typ is DictType:
-        if maxlevels and level >= maxlevels:
-            s = "{...}"
-            readable = 0
-        else:
-            items = object.items()
-            k, v = items[0]
+    startchar, endchar = {ListType:  "[]",
+                          TupleType: "()",
+                          DictType:  "{}"}[typ]
+    if maxlevels and level > maxlevels:
+        with_commas = "..."
+        readable = 0
+
+    elif typ is DictType:
+        components = []
+        for k, v in object.iteritems():
             krepr, kreadable, krecur = _safe_repr(k, context, maxlevels,
                                                   level)
             vrepr, vreadable, vrecur = _safe_repr(v, context, maxlevels,
                                                   level)
+            components.append("%s: %s" % (krepr, vrepr))
             readable = readable and kreadable and vreadable
             recursive = recursive or krecur or vrecur
-            s = "{%s: %s" % (krepr, vrepr)
-            for k, v in items[1:]:
-                krepr, kreadable, krecur = _safe_repr(k, context, maxlevels,
-                                                      level)
-                vrepr, vreadable, vrecur = _safe_repr(v, context, maxlevels,
-                                                      level)
-                readable = readable and kreadable and vreadable
-                recursive = recursive or krecur or vrecur
-                s = "%s, %s: %s" % (s, krepr, vrepr)
-            s = s + "}"
-    else:
-        s, term = (typ is ListType) and ('[', ']') or ('(', ')')
-        if maxlevels and level >= maxlevels:
-            s = s + "..."
-            readable = 0
-        else:
+        with_commas = ", ".join(components)
+
+    else: # list or tuple
+        assert typ in (ListType, TupleType)
+        components = []
+        for element in object:
             subrepr, subreadable, subrecur = _safe_repr(
-                object[0], context, maxlevels, level)
+                  element, context, maxlevels, level)
+            components.append(subrepr)
             readable = readable and subreadable
             recursive = recursive or subrecur
-            s = s + subrepr
-            tail = object[1:]
-            if not tail:
-                if typ is TupleType:
-                    s = s + ','
-            for ent in tail:
-                subrepr, subreadable, subrecur = _safe_repr(
-                    ent, context, maxlevels, level)
-                readable = readable and subreadable
-                recursive = recursive or subrecur
-                s = "%s, %s" % (s, subrepr)
-        s = s + term
+        if len(components) == 1 and typ is TupleType:
+            components[0] += ","
+        with_commas = ", ".join(components)
+
+    s = "%s%s%s" % (startchar, with_commas, endchar)
     del context[objid]
     return s, readable and not recursive, recursive
 
-
 class _Recursion:
     # represent a recursive relationship; really only used for the __repr__()
     # method...
diff --git a/Lib/test/test_pprint.py b/Lib/test/test_pprint.py
index f6a128d..25d806d 100644
--- a/Lib/test/test_pprint.py
+++ b/Lib/test/test_pprint.py
@@ -12,7 +12,7 @@ class QueryTestCase(unittest.TestCase):
         self.a[-12] = self.b
 
     def test_basic(self):
-        # Verify that .isrecursive() and .isreadable() work.
+        # Verify that .isrecursive() and .isreadable() work w/o recursion.
         verify = self.assert_
         for safe in (2, 2.0, 2j, "abc", [3], (2,2), {3: 3}, u"yaddayadda",
                      self.a, self.b):
@@ -22,6 +22,7 @@ class QueryTestCase(unittest.TestCase):
                    "expected isreadable for " + `safe`)
 
     def test_knotted(self):
+        # Verify that .isrecursive() and .isreadable() work w/ recursion.
         # Tie a knot.
         self.b[67] = self.a
         # Messy dict.
@@ -54,5 +55,20 @@ class QueryTestCase(unittest.TestCase):
             verify(not pprint.isreadable(unreadable),
                    "expected not isreadable for " + `unreadable`)
 
+    def test_same_as_repr(self):
+        "Simple objects and small containers that should be same as repr()."
+        verify = self.assert_
+        for simple in (0, 0L, 0+0j, 0.0, "", u"", (), [], {}, verify, pprint,
+                       -6, -6L, -6-6j, -1.5, "x", u"x", (3,), [3], {3: 6},
+                       (1,2), [3,4], {5: 6, 7: 8},
+                       {"xy\tab\n": (3,), 5: [[]], (): {}},
+                       range(10, -11, -1)
+                      ):
+            native = repr(simple)
+            for function in "pformat", "saferepr":
+                f = getattr(pprint, function)
+                got = f(simple)
+                verify(native == got, "expected %s got %s from pprint.%s" %
+                                      (native, got, function))
 
 test_support.run_unittest(QueryTestCase)