1#  Author:      Fred L. Drake, Jr.
2#               fdrake@acm.org
3#
4#  This is a simple little module I wrote to make life easier.  I didn't
5#  see anything quite like it in the library, though I may have overlooked
6#  something.  I wrote this when I was trying to read some heavily nested
7#  tuples with fairly non-descriptive content.  This is modeled very much
8#  after Lisp/Scheme - style pretty-printing of lists.  If you find it
9#  useful, thank small children who sleep at night.
10
11"""Support to pretty-print lists, tuples, & dictionaries recursively.
12
13Very simple, but useful, especially in debugging data structures.
14
15Classes
16-------
17
18PrettyPrinter()
19    Handle pretty-printing operations onto a stream using a configured
20    set of formatting parameters.
21
22Functions
23---------
24
25pformat()
26    Format a Python object into a pretty-printed representation.
27
28pprint()
29    Pretty-print a Python object to a stream [default is sys.stdout].
30
31saferepr()
32    Generate a 'standard' repr()-like value, but protect against recursive
33    data structures.
34
35"""
36
37import collections as _collections
38import re
39import sys as _sys
40import types as _types
41from io import StringIO as _StringIO
42
43__all__ = ["pprint","pformat","isreadable","isrecursive","saferepr",
44           "PrettyPrinter", "pp"]
45
46
47def pprint(object, stream=None, indent=1, width=80, depth=None, *,
48           compact=False, sort_dicts=True):
49    """Pretty-print a Python object to a stream [default is sys.stdout]."""
50    printer = PrettyPrinter(
51        stream=stream, indent=indent, width=width, depth=depth,
52        compact=compact, sort_dicts=sort_dicts)
53    printer.pprint(object)
54
55def pformat(object, indent=1, width=80, depth=None, *,
56            compact=False, sort_dicts=True):
57    """Format a Python object into a pretty-printed representation."""
58    return PrettyPrinter(indent=indent, width=width, depth=depth,
59                         compact=compact, sort_dicts=sort_dicts).pformat(object)
60
61def pp(object, *args, sort_dicts=False, **kwargs):
62    """Pretty-print a Python object"""
63    pprint(object, *args, sort_dicts=sort_dicts, **kwargs)
64
65def saferepr(object):
66    """Version of repr() which can handle recursive data structures."""
67    return _safe_repr(object, {}, None, 0, True)[0]
68
69def isreadable(object):
70    """Determine if saferepr(object) is readable by eval()."""
71    return _safe_repr(object, {}, None, 0, True)[1]
72
73def isrecursive(object):
74    """Determine if object requires a recursive representation."""
75    return _safe_repr(object, {}, None, 0, True)[2]
76
77class _safe_key:
78    """Helper function for key functions when sorting unorderable objects.
79
80    The wrapped-object will fallback to a Py2.x style comparison for
81    unorderable types (sorting first comparing the type name and then by
82    the obj ids).  Does not work recursively, so dict.items() must have
83    _safe_key applied to both the key and the value.
84
85    """
86
87    __slots__ = ['obj']
88
89    def __init__(self, obj):
90        self.obj = obj
91
92    def __lt__(self, other):
93        try:
94            return self.obj < other.obj
95        except TypeError:
96            return ((str(type(self.obj)), id(self.obj)) < \
97                    (str(type(other.obj)), id(other.obj)))
98
99def _safe_tuple(t):
100    "Helper function for comparing 2-tuples"
101    return _safe_key(t[0]), _safe_key(t[1])
102
103class PrettyPrinter:
104    def __init__(self, indent=1, width=80, depth=None, stream=None, *,
105                 compact=False, sort_dicts=True):
106        """Handle pretty printing operations onto a stream using a set of
107        configured parameters.
108
109        indent
110            Number of spaces to indent for each level of nesting.
111
112        width
113            Attempted maximum number of columns in the output.
114
115        depth
116            The maximum depth to print out nested structures.
117
118        stream
119            The desired output stream.  If omitted (or false), the standard
120            output stream available at construction will be used.
121
122        compact
123            If true, several items will be combined in one line.
124
125        sort_dicts
126            If true, dict keys are sorted.
127
128        """
129        indent = int(indent)
130        width = int(width)
131        if indent < 0:
132            raise ValueError('indent must be >= 0')
133        if depth is not None and depth <= 0:
134            raise ValueError('depth must be > 0')
135        if not width:
136            raise ValueError('width must be != 0')
137        self._depth = depth
138        self._indent_per_level = indent
139        self._width = width
140        if stream is not None:
141            self._stream = stream
142        else:
143            self._stream = _sys.stdout
144        self._compact = bool(compact)
145        self._sort_dicts = sort_dicts
146
147    def pprint(self, object):
148        self._format(object, self._stream, 0, 0, {}, 0)
149        self._stream.write("\n")
150
151    def pformat(self, object):
152        sio = _StringIO()
153        self._format(object, sio, 0, 0, {}, 0)
154        return sio.getvalue()
155
156    def isrecursive(self, object):
157        return self.format(object, {}, 0, 0)[2]
158
159    def isreadable(self, object):
160        s, readable, recursive = self.format(object, {}, 0, 0)
161        return readable and not recursive
162
163    def _format(self, object, stream, indent, allowance, context, level):
164        objid = id(object)
165        if objid in context:
166            stream.write(_recursion(object))
167            self._recursive = True
168            self._readable = False
169            return
170        rep = self._repr(object, context, level)
171        max_width = self._width - indent - allowance
172        if len(rep) > max_width:
173            p = self._dispatch.get(type(object).__repr__, None)
174            if p is not None:
175                context[objid] = 1
176                p(self, object, stream, indent, allowance, context, level + 1)
177                del context[objid]
178                return
179            elif isinstance(object, dict):
180                context[objid] = 1
181                self._pprint_dict(object, stream, indent, allowance,
182                                  context, level + 1)
183                del context[objid]
184                return
185        stream.write(rep)
186
187    _dispatch = {}
188
189    def _pprint_dict(self, object, stream, indent, allowance, context, level):
190        write = stream.write
191        write('{')
192        if self._indent_per_level > 1:
193            write((self._indent_per_level - 1) * ' ')
194        length = len(object)
195        if length:
196            if self._sort_dicts:
197                items = sorted(object.items(), key=_safe_tuple)
198            else:
199                items = object.items()
200            self._format_dict_items(items, stream, indent, allowance + 1,
201                                    context, level)
202        write('}')
203
204    _dispatch[dict.__repr__] = _pprint_dict
205
206    def _pprint_ordered_dict(self, object, stream, indent, allowance, context, level):
207        if not len(object):
208            stream.write(repr(object))
209            return
210        cls = object.__class__
211        stream.write(cls.__name__ + '(')
212        self._format(list(object.items()), stream,
213                     indent + len(cls.__name__) + 1, allowance + 1,
214                     context, level)
215        stream.write(')')
216
217    _dispatch[_collections.OrderedDict.__repr__] = _pprint_ordered_dict
218
219    def _pprint_list(self, object, stream, indent, allowance, context, level):
220        stream.write('[')
221        self._format_items(object, stream, indent, allowance + 1,
222                           context, level)
223        stream.write(']')
224
225    _dispatch[list.__repr__] = _pprint_list
226
227    def _pprint_tuple(self, object, stream, indent, allowance, context, level):
228        stream.write('(')
229        endchar = ',)' if len(object) == 1 else ')'
230        self._format_items(object, stream, indent, allowance + len(endchar),
231                           context, level)
232        stream.write(endchar)
233
234    _dispatch[tuple.__repr__] = _pprint_tuple
235
236    def _pprint_set(self, object, stream, indent, allowance, context, level):
237        if not len(object):
238            stream.write(repr(object))
239            return
240        typ = object.__class__
241        if typ is set:
242            stream.write('{')
243            endchar = '}'
244        else:
245            stream.write(typ.__name__ + '({')
246            endchar = '})'
247            indent += len(typ.__name__) + 1
248        object = sorted(object, key=_safe_key)
249        self._format_items(object, stream, indent, allowance + len(endchar),
250                           context, level)
251        stream.write(endchar)
252
253    _dispatch[set.__repr__] = _pprint_set
254    _dispatch[frozenset.__repr__] = _pprint_set
255
256    def _pprint_str(self, object, stream, indent, allowance, context, level):
257        write = stream.write
258        if not len(object):
259            write(repr(object))
260            return
261        chunks = []
262        lines = object.splitlines(True)
263        if level == 1:
264            indent += 1
265            allowance += 1
266        max_width1 = max_width = self._width - indent
267        for i, line in enumerate(lines):
268            rep = repr(line)
269            if i == len(lines) - 1:
270                max_width1 -= allowance
271            if len(rep) <= max_width1:
272                chunks.append(rep)
273            else:
274                # A list of alternating (non-space, space) strings
275                parts = re.findall(r'\S*\s*', line)
276                assert parts
277                assert not parts[-1]
278                parts.pop()  # drop empty last part
279                max_width2 = max_width
280                current = ''
281                for j, part in enumerate(parts):
282                    candidate = current + part
283                    if j == len(parts) - 1 and i == len(lines) - 1:
284                        max_width2 -= allowance
285                    if len(repr(candidate)) > max_width2:
286                        if current:
287                            chunks.append(repr(current))
288                        current = part
289                    else:
290                        current = candidate
291                if current:
292                    chunks.append(repr(current))
293        if len(chunks) == 1:
294            write(rep)
295            return
296        if level == 1:
297            write('(')
298        for i, rep in enumerate(chunks):
299            if i > 0:
300                write('\n' + ' '*indent)
301            write(rep)
302        if level == 1:
303            write(')')
304
305    _dispatch[str.__repr__] = _pprint_str
306
307    def _pprint_bytes(self, object, stream, indent, allowance, context, level):
308        write = stream.write
309        if len(object) <= 4:
310            write(repr(object))
311            return
312        parens = level == 1
313        if parens:
314            indent += 1
315            allowance += 1
316            write('(')
317        delim = ''
318        for rep in _wrap_bytes_repr(object, self._width - indent, allowance):
319            write(delim)
320            write(rep)
321            if not delim:
322                delim = '\n' + ' '*indent
323        if parens:
324            write(')')
325
326    _dispatch[bytes.__repr__] = _pprint_bytes
327
328    def _pprint_bytearray(self, object, stream, indent, allowance, context, level):
329        write = stream.write
330        write('bytearray(')
331        self._pprint_bytes(bytes(object), stream, indent + 10,
332                           allowance + 1, context, level + 1)
333        write(')')
334
335    _dispatch[bytearray.__repr__] = _pprint_bytearray
336
337    def _pprint_mappingproxy(self, object, stream, indent, allowance, context, level):
338        stream.write('mappingproxy(')
339        self._format(object.copy(), stream, indent + 13, allowance + 1,
340                     context, level)
341        stream.write(')')
342
343    _dispatch[_types.MappingProxyType.__repr__] = _pprint_mappingproxy
344
345    def _format_dict_items(self, items, stream, indent, allowance, context,
346                           level):
347        write = stream.write
348        indent += self._indent_per_level
349        delimnl = ',\n' + ' ' * indent
350        last_index = len(items) - 1
351        for i, (key, ent) in enumerate(items):
352            last = i == last_index
353            rep = self._repr(key, context, level)
354            write(rep)
355            write(': ')
356            self._format(ent, stream, indent + len(rep) + 2,
357                         allowance if last else 1,
358                         context, level)
359            if not last:
360                write(delimnl)
361
362    def _format_items(self, items, stream, indent, allowance, context, level):
363        write = stream.write
364        indent += self._indent_per_level
365        if self._indent_per_level > 1:
366            write((self._indent_per_level - 1) * ' ')
367        delimnl = ',\n' + ' ' * indent
368        delim = ''
369        width = max_width = self._width - indent + 1
370        it = iter(items)
371        try:
372            next_ent = next(it)
373        except StopIteration:
374            return
375        last = False
376        while not last:
377            ent = next_ent
378            try:
379                next_ent = next(it)
380            except StopIteration:
381                last = True
382                max_width -= allowance
383                width -= allowance
384            if self._compact:
385                rep = self._repr(ent, context, level)
386                w = len(rep) + 2
387                if width < w:
388                    width = max_width
389                    if delim:
390                        delim = delimnl
391                if width >= w:
392                    width -= w
393                    write(delim)
394                    delim = ', '
395                    write(rep)
396                    continue
397            write(delim)
398            delim = delimnl
399            self._format(ent, stream, indent,
400                         allowance if last else 1,
401                         context, level)
402
403    def _repr(self, object, context, level):
404        repr, readable, recursive = self.format(object, context.copy(),
405                                                self._depth, level)
406        if not readable:
407            self._readable = False
408        if recursive:
409            self._recursive = True
410        return repr
411
412    def format(self, object, context, maxlevels, level):
413        """Format object for a specific context, returning a string
414        and flags indicating whether the representation is 'readable'
415        and whether the object represents a recursive construct.
416        """
417        return _safe_repr(object, context, maxlevels, level, self._sort_dicts)
418
419    def _pprint_default_dict(self, object, stream, indent, allowance, context, level):
420        if not len(object):
421            stream.write(repr(object))
422            return
423        rdf = self._repr(object.default_factory, context, level)
424        cls = object.__class__
425        indent += len(cls.__name__) + 1
426        stream.write('%s(%s,\n%s' % (cls.__name__, rdf, ' ' * indent))
427        self._pprint_dict(object, stream, indent, allowance + 1, context, level)
428        stream.write(')')
429
430    _dispatch[_collections.defaultdict.__repr__] = _pprint_default_dict
431
432    def _pprint_counter(self, object, stream, indent, allowance, context, level):
433        if not len(object):
434            stream.write(repr(object))
435            return
436        cls = object.__class__
437        stream.write(cls.__name__ + '({')
438        if self._indent_per_level > 1:
439            stream.write((self._indent_per_level - 1) * ' ')
440        items = object.most_common()
441        self._format_dict_items(items, stream,
442                                indent + len(cls.__name__) + 1, allowance + 2,
443                                context, level)
444        stream.write('})')
445
446    _dispatch[_collections.Counter.__repr__] = _pprint_counter
447
448    def _pprint_chain_map(self, object, stream, indent, allowance, context, level):
449        if not len(object.maps):
450            stream.write(repr(object))
451            return
452        cls = object.__class__
453        stream.write(cls.__name__ + '(')
454        indent += len(cls.__name__) + 1
455        for i, m in enumerate(object.maps):
456            if i == len(object.maps) - 1:
457                self._format(m, stream, indent, allowance + 1, context, level)
458                stream.write(')')
459            else:
460                self._format(m, stream, indent, 1, context, level)
461                stream.write(',\n' + ' ' * indent)
462
463    _dispatch[_collections.ChainMap.__repr__] = _pprint_chain_map
464
465    def _pprint_deque(self, object, stream, indent, allowance, context, level):
466        if not len(object):
467            stream.write(repr(object))
468            return
469        cls = object.__class__
470        stream.write(cls.__name__ + '(')
471        indent += len(cls.__name__) + 1
472        stream.write('[')
473        if object.maxlen is None:
474            self._format_items(object, stream, indent, allowance + 2,
475                               context, level)
476            stream.write('])')
477        else:
478            self._format_items(object, stream, indent, 2,
479                               context, level)
480            rml = self._repr(object.maxlen, context, level)
481            stream.write('],\n%smaxlen=%s)' % (' ' * indent, rml))
482
483    _dispatch[_collections.deque.__repr__] = _pprint_deque
484
485    def _pprint_user_dict(self, object, stream, indent, allowance, context, level):
486        self._format(object.data, stream, indent, allowance, context, level - 1)
487
488    _dispatch[_collections.UserDict.__repr__] = _pprint_user_dict
489
490    def _pprint_user_list(self, object, stream, indent, allowance, context, level):
491        self._format(object.data, stream, indent, allowance, context, level - 1)
492
493    _dispatch[_collections.UserList.__repr__] = _pprint_user_list
494
495    def _pprint_user_string(self, object, stream, indent, allowance, context, level):
496        self._format(object.data, stream, indent, allowance, context, level - 1)
497
498    _dispatch[_collections.UserString.__repr__] = _pprint_user_string
499
500# Return triple (repr_string, isreadable, isrecursive).
501
502def _safe_repr(object, context, maxlevels, level, sort_dicts):
503    typ = type(object)
504    if typ in _builtin_scalars:
505        return repr(object), True, False
506
507    r = getattr(typ, "__repr__", None)
508    if issubclass(typ, dict) and r is dict.__repr__:
509        if not object:
510            return "{}", True, False
511        objid = id(object)
512        if maxlevels and level >= maxlevels:
513            return "{...}", False, objid in context
514        if objid in context:
515            return _recursion(object), False, True
516        context[objid] = 1
517        readable = True
518        recursive = False
519        components = []
520        append = components.append
521        level += 1
522        if sort_dicts:
523            items = sorted(object.items(), key=_safe_tuple)
524        else:
525            items = object.items()
526        for k, v in items:
527            krepr, kreadable, krecur = _safe_repr(k, context, maxlevels, level, sort_dicts)
528            vrepr, vreadable, vrecur = _safe_repr(v, context, maxlevels, level, sort_dicts)
529            append("%s: %s" % (krepr, vrepr))
530            readable = readable and kreadable and vreadable
531            if krecur or vrecur:
532                recursive = True
533        del context[objid]
534        return "{%s}" % ", ".join(components), readable, recursive
535
536    if (issubclass(typ, list) and r is list.__repr__) or \
537       (issubclass(typ, tuple) and r is tuple.__repr__):
538        if issubclass(typ, list):
539            if not object:
540                return "[]", True, False
541            format = "[%s]"
542        elif len(object) == 1:
543            format = "(%s,)"
544        else:
545            if not object:
546                return "()", True, False
547            format = "(%s)"
548        objid = id(object)
549        if maxlevels and level >= maxlevels:
550            return format % "...", False, objid in context
551        if objid in context:
552            return _recursion(object), False, True
553        context[objid] = 1
554        readable = True
555        recursive = False
556        components = []
557        append = components.append
558        level += 1
559        for o in object:
560            orepr, oreadable, orecur = _safe_repr(o, context, maxlevels, level, sort_dicts)
561            append(orepr)
562            if not oreadable:
563                readable = False
564            if orecur:
565                recursive = True
566        del context[objid]
567        return format % ", ".join(components), readable, recursive
568
569    rep = repr(object)
570    return rep, (rep and not rep.startswith('<')), False
571
572_builtin_scalars = frozenset({str, bytes, bytearray, int, float, complex,
573                              bool, type(None)})
574
575def _recursion(object):
576    return ("<Recursion on %s with id=%s>"
577            % (type(object).__name__, id(object)))
578
579
580def _perfcheck(object=None):
581    import time
582    if object is None:
583        object = [("string", (1, 2), [3, 4], {5: 6, 7: 8})] * 100000
584    p = PrettyPrinter()
585    t1 = time.perf_counter()
586    _safe_repr(object, {}, None, 0, True)
587    t2 = time.perf_counter()
588    p.pformat(object)
589    t3 = time.perf_counter()
590    print("_safe_repr:", t2 - t1)
591    print("pformat:", t3 - t2)
592
593def _wrap_bytes_repr(object, width, allowance):
594    current = b''
595    last = len(object) // 4 * 4
596    for i in range(0, len(object), 4):
597        part = object[i: i+4]
598        candidate = current + part
599        if i == last:
600            width -= allowance
601        if len(repr(candidate)) > width:
602            if current:
603                yield repr(current)
604            current = part
605        else:
606            current = candidate
607    if current:
608        yield repr(current)
609
610if __name__ == "__main__":
611    _perfcheck()
612