pprint: implement depth recursion

Author: mikebentley15

python-stdlib/pprint/pprint.py

@@ -1,6 +1,10 @@
-from io import StringIO as _StringIO
-from sys import stdout as _stdout
-from machine import const as _const
+import io as _io
+import sys as _sys
+
+try:
+    from micropython import const as _const
+except ImportError:
+    _const = lambda x: x
 
 
 class PrettyPrinter:
@@ -10,27 +14,34 @@ class PrettyPrinter:
     call.  For parameter descriptions, look at pprint() documentation.
     """
 
-    def __init__(self, indent=1, stream=None, sort_dicts=True):
-        self.indent = indent
-        self.stream = stream if stream else _stdout
-        self.sort_dicts = sort_dicts
+    def __init__(self, indent=1, depth=1, stream=None, sort_dicts=True):
+        self._indent = indent
+        self._depth = depth
+        self._stream = stream if stream else _sys.stdout
+        self._sort_dicts = sort_dicts
 
     def pformat(self, obj):
         """Call pformat() function with cached constructor values"""
-        return pformat(obj, indent=self.indent, sort_dicts=self.sort_dicts)
+        return pformat(obj, indent=self._indent, depth=self._depth, sort_dicts=self._sort_dicts)
 
     def pprint(self, obj):
         """Call pprint() function with cached constructor values"""
-        pprint(obj, stream=self.stream, indent=self.indent, sort_dicts=self.sort_dicts)
+        pprint(
+            obj,
+            stream=self._stream,
+            indent=self._indent,
+            depth=self._depth,
+            sort_dicts=self._sort_dicts,
+        )
 
 
-def pformat(obj, indent=1, sort_dicts=True):
+def pformat(obj, indent=1, depth=1, sort_dicts=True):
     """Simply calls pprint() but capturing prints into a string and returning.
 
     Does not add a newline to the end like pprint() does.
     """
-    stream = _StringIO()
-    _pprint_impl(obj, stream=stream, indent=indent, sort_dicts=sort_dicts)
+    stream = _io.StringIO()
+    _pprint_impl(obj, stream, 0, sort_dicts, depth, 0, indent)
     return stream.getvalue()
 
 
@@ -39,7 +50,7 @@ def pp(obj, *args, sort_dicts=False, **kwargs):
     pprint(obj, *args, sort_dicts=sort_dicts, **kwargs)
 
 
-def pprint(obj, stream=None, indent=1, sort_dicts=True):
+def pprint(obj, stream=None, indent=1, depth=1, sort_dicts=True):
     """Simple implementation of a pretty-printer.
 
     For simplicity, this does not recurse down.  It does not produce the same
@@ -49,13 +60,26 @@ def pprint(obj, stream=None, indent=1, sort_dicts=True):
     Only expands containers.  Currently supported containers are set, list,
     dict, and tuple.
 
+    The depth parameter defaults to 1 which is different from the CPython
+    implementation which defaults to None, meaning infinite recursion.  This
+    implementation is significantly simpler and does not have safety checks
+    against recursive structures (i.e., structures that contain a copy of
+    themselves somewhere) or costly overexpansion.  This simplicity gains speed
+    and memory efficiency at the potential cost of expanding too far.  Increase
+    the depth parameter cautiously.
+
     stream
         The desired output stream.  If omitted (or False), the standard output
         stream will be used.
 
     indent
         Number of space to indent for each level of nesting.
 
+    depth
+        The maximum depth to print out nested structures.  Set to None for
+        infinite depth.  Warning: be careful of infinite depth as infinite
+        recursion and overexpansion checks are not performed.
+
     sort_dicts
         If true, dict keys are sorted.  You must ensure that the dicts
         encountered can be sorted against each other (for example, int cannot
@@ -65,8 +89,8 @@ def pprint(obj, stream=None, indent=1, sort_dicts=True):
     mismatch
     """
     if not stream:
-        stream = _stdout
-    _pprint_impl(obj, stream, indent, sort_dicts)
+        stream = _sys.stdout
+    _pprint_impl(obj, stream, 0, sort_dicts, depth, 0, indent)
     stream.write("\n")  # end in a newline
 
 
@@ -85,16 +109,36 @@ def ppdir(obj, *args, hidden=False, **kwargs):
 # Private implementation
 
 
-def _pprint_impl(obj, stream, indent, sort_dicts):
+def _pprint_impl(obj, stream, indent, sort_dicts, maxlevel, level, indent_per_level):
     type_code = _container_type(obj)
-    if type_code != _NONCONTAINER and len(obj) > 1:
+    if (maxlevel is not None and level >= maxlevel) or type_code == _NONCONTAINER or not obj:
+        # directly print noncontainers, empty containers, or if we have reached
+        # the max level.
+        stream.write(repr(obj))
+    else:
+        # expand the container
         if type_code & _SIMPLE:
-            _pprint_simple_container(obj, type_code, stream, indent, sort_dicts)
+            _pprint_simple_container(
+                obj,
+                type_code,
+                stream,
+                indent + indent_per_level,
+                sort_dicts,
+                maxlevel,
+                level + 1,
+                indent_per_level,
+            )
         else:  # type_code & _DICT:
-            _pprint_dict(obj, type_code, stream, indent, sort_dicts)
-    else:
-        # directly print noncontainers or containers with one or fewer elements
-        stream.write(repr(obj))
+            _pprint_dict(
+                obj,
+                type_code,
+                stream,
+                indent + indent_per_level,
+                sort_dicts,
+                maxlevel,
+                level + 1,
+                indent_per_level,
+            )
 
 
 _NONCONTAINER = _const(0)
@@ -118,37 +162,55 @@ def _container_type(obj):
     return _NONCONTAINER
 
 
-def _pprint_simple_container(obj, type_code, stream, indent, sort_dicts):
+def _pprint_simple_container(
+    obj, type_code, stream, indent, sort_dicts, maxlevel, level, indent_per_level
+):
     if type_code & _LIST:
         chars = "[]"
     elif type_code & _TUPLE:
         chars = "()"
     else:  # type_code & _SET:
         chars = "{}"
 
-    stream.write(chars[0] + " " * (indent - 1))
+    indent_string = " " * indent
+
+    stream.write(chars[0])
+    stream.write(" " * (indent_per_level - 1))
     first = True
     for item in obj:
         if not first:
-            stream.write(",\n" + " " * indent)
+            stream.write(",\n")
+            stream.write(indent_string)
         first = False
-        # TODO: would recurse here
-        stream.write(repr(item))
+        # recurse here
+        _pprint_impl(item, stream, indent, sort_dicts, maxlevel, level, indent_per_level)
+    if type_code & _TUPLE and len(obj) == 1:
+        stream.write(",")
     stream.write(chars[1])
 
 
-def _pprint_dict(obj, type_code, stream, indent, sort_dicts):
-    stream.write("{" + " " * (indent - 1))
+def _pprint_dict(obj, type_code, stream, indent, sort_dicts, maxlevel, level, indent_per_level):
+    indent_string = " " * indent
+    stream.write("{")
+    stream.write(" " * (indent_per_level - 1))
     first = True
     items = obj.items()
     if sort_dicts:
         items = sorted(items, key=lambda x: x[0])
     for key, val in items:
         if not first:
-            stream.write(",\n" + " " * indent)
+            stream.write(",\n")
+            stream.write(indent_string)
         first = False
-        # TODO: would recurse here
-        stream.write(repr(key))
+        # python's implementation does not multiline keys, then offsets the
+        # value by the key's length.  They still do safety checks such as
+        # recursiveness and checking that it can be printed within the
+        # specified width, which we skip in this implementation.
+        keystr = repr(key)
+        stream.write(keystr)
         stream.write(": ")
-        stream.write(repr(val))
+        # recurse on the value, lined up to wrap in front of the colon.
+        _pprint_impl(
+            val, stream, indent + len(keystr) + 2, sort_dicts, maxlevel, level, indent_per_level
+        )
     stream.write("}")