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 _pprint_simplenamespace(self, object, stream, indent, allowance, context, level):
346        if type(object) is _types.SimpleNamespace:
347            # The SimpleNamespace repr is "namespace" instead of the class
348            # name, so we do the same here. For subclasses; use the class name.
349            cls_name = 'namespace'
350        else:
351            cls_name = object.__class__.__name__
352        indent += len(cls_name) + 1
353        delimnl = ',\n' + ' ' * indent
354        items = object.__dict__.items()
355        last_index = len(items) - 1
356
357        stream.write(cls_name + '(')
358        for i, (key, ent) in enumerate(items):
359            stream.write(key)
360            stream.write('=')
361
362            last = i == last_index
363            self._format(ent, stream, indent + len(key) + 1,
364                         allowance if last else 1,
365                         context, level)
366            if not last:
367                stream.write(delimnl)
368        stream.write(')')
369
370    _dispatch[_types.SimpleNamespace.__repr__] = _pprint_simplenamespace
371
372    def _format_dict_items(self, items, stream, indent, allowance, context,
373                           level):
374        write = stream.write
375        indent += self._indent_per_level
376        delimnl = ',\n' + ' ' * indent
377        last_index = len(items) - 1
378        for i, (key, ent) in enumerate(items):
379            last = i == last_index
380            rep = self._repr(key, context, level)
381            write(rep)
382            write(': ')
383            self._format(ent, stream, indent + len(rep) + 2,
384                         allowance if last else 1,
385                         context, level)
386            if not last:
387                write(delimnl)
388
389    def _format_items(self, items, stream, indent, allowance, context, level):
390        write = stream.write
391        indent += self._indent_per_level
392        if self._indent_per_level > 1:
393            write((self._indent_per_level - 1) * ' ')
394        delimnl = ',\n' + ' ' * indent
395        delim = ''
396        width = max_width = self._width - indent + 1
397        it = iter(items)
398        try:
399            next_ent = next(it)
400        except StopIteration:
401            return
402        last = False
403        while not last:
404            ent = next_ent
405            try:
406                next_ent = next(it)
407            except StopIteration:
408                last = True
409                max_width -= allowance
410                width -= allowance
411            if self._compact:
412                rep = self._repr(ent, context, level)
413                w = len(rep) + 2
414                if width < w:
415                    width = max_width
416                    if delim:
417                        delim = delimnl
418                if width >= w:
419                    width -= w
420                    write(delim)
421                    delim = ', '
422                    write(rep)
423                    continue
424            write(delim)
425            delim = delimnl
426            self._format(ent, stream, indent,
427                         allowance if last else 1,
428                         context, level)
429
430    def _repr(self, object, context, level):
431        repr, readable, recursive = self.format(object, context.copy(),
432                                                self._depth, level)
433        if not readable:
434            self._readable = False
435        if recursive:
436            self._recursive = True
437        return repr
438
439    def format(self, object, context, maxlevels, level):
440        """Format object for a specific context, returning a string
441        and flags indicating whether the representation is 'readable'
442        and whether the object represents a recursive construct.
443        """
444        return _safe_repr(object, context, maxlevels, level, self._sort_dicts)
445
446    def _pprint_default_dict(self, object, stream, indent, allowance, context, level):
447        if not len(object):
448            stream.write(repr(object))
449            return
450        rdf = self._repr(object.default_factory, context, level)
451        cls = object.__class__
452        indent += len(cls.__name__) + 1
453        stream.write('%s(%s,\n%s' % (cls.__name__, rdf, ' ' * indent))
454        self._pprint_dict(object, stream, indent, allowance + 1, context, level)
455        stream.write(')')
456
457    _dispatch[_collections.defaultdict.__repr__] = _pprint_default_dict
458
459    def _pprint_counter(self, object, stream, indent, allowance, context, level):
460        if not len(object):
461            stream.write(repr(object))
462            return
463        cls = object.__class__
464        stream.write(cls.__name__ + '({')
465        if self._indent_per_level > 1:
466            stream.write((self._indent_per_level - 1) * ' ')
467        items = object.most_common()
468        self._format_dict_items(items, stream,
469                                indent + len(cls.__name__) + 1, allowance + 2,
470                                context, level)
471        stream.write('})')
472
473    _dispatch[_collections.Counter.__repr__] = _pprint_counter
474
475    def _pprint_chain_map(self, object, stream, indent, allowance, context, level):
476        if not len(object.maps):
477            stream.write(repr(object))
478            return
479        cls = object.__class__
480        stream.write(cls.__name__ + '(')
481        indent += len(cls.__name__) + 1
482        for i, m in enumerate(object.maps):
483            if i == len(object.maps) - 1:
484                self._format(m, stream, indent, allowance + 1, context, level)
485                stream.write(')')
486            else:
487                self._format(m, stream, indent, 1, context, level)
488                stream.write(',\n' + ' ' * indent)
489
490    _dispatch[_collections.ChainMap.__repr__] = _pprint_chain_map
491
492    def _pprint_deque(self, object, stream, indent, allowance, context, level):
493        if not len(object):
494            stream.write(repr(object))
495            return
496        cls = object.__class__
497        stream.write(cls.__name__ + '(')
498        indent += len(cls.__name__) + 1
499        stream.write('[')
500        if object.maxlen is None:
501            self._format_items(object, stream, indent, allowance + 2,
502                               context, level)
503            stream.write('])')
504        else:
505            self._format_items(object, stream, indent, 2,
506                               context, level)
507            rml = self._repr(object.maxlen, context, level)
508            stream.write('],\n%smaxlen=%s)' % (' ' * indent, rml))
509
510    _dispatch[_collections.deque.__repr__] = _pprint_deque
511
512    def _pprint_user_dict(self, object, stream, indent, allowance, context, level):
513        self._format(object.data, stream, indent, allowance, context, level - 1)
514
515    _dispatch[_collections.UserDict.__repr__] = _pprint_user_dict
516
517    def _pprint_user_list(self, object, stream, indent, allowance, context, level):
518        self._format(object.data, stream, indent, allowance, context, level - 1)
519
520    _dispatch[_collections.UserList.__repr__] = _pprint_user_list
521
522    def _pprint_user_string(self, object, stream, indent, allowance, context, level):
523        self._format(object.data, stream, indent, allowance, context, level - 1)
524
525    _dispatch[_collections.UserString.__repr__] = _pprint_user_string
526
527# Return triple (repr_string, isreadable, isrecursive).
528
529def _safe_repr(object, context, maxlevels, level, sort_dicts):
530    typ = type(object)
531    if typ in _builtin_scalars:
532        return repr(object), True, False
533
534    r = getattr(typ, "__repr__", None)
535    if issubclass(typ, dict) and r is dict.__repr__:
536        if not object:
537            return "{}", True, False
538        objid = id(object)
539        if maxlevels and level >= maxlevels:
540            return "{...}", False, objid in context
541        if objid in context:
542            return _recursion(object), False, True
543        context[objid] = 1
544        readable = True
545        recursive = False
546        components = []
547        append = components.append
548        level += 1
549        if sort_dicts:
550            items = sorted(object.items(), key=_safe_tuple)
551        else:
552            items = object.items()
553        for k, v in items:
554            krepr, kreadable, krecur = _safe_repr(k, context, maxlevels, level, sort_dicts)
555            vrepr, vreadable, vrecur = _safe_repr(v, context, maxlevels, level, sort_dicts)
556            append("%s: %s" % (krepr, vrepr))
557            readable = readable and kreadable and vreadable
558            if krecur or vrecur:
559                recursive = True
560        del context[objid]
561        return "{%s}" % ", ".join(components), readable, recursive
562
563    if (issubclass(typ, list) and r is list.__repr__) or \
564       (issubclass(typ, tuple) and r is tuple.__repr__):
565        if issubclass(typ, list):
566            if not object:
567                return "[]", True, False
568            format = "[%s]"
569        elif len(object) == 1:
570            format = "(%s,)"
571        else:
572            if not object:
573                return "()", True, False
574            format = "(%s)"
575        objid = id(object)
576        if maxlevels and level >= maxlevels:
577            return format % "...", False, objid in context
578        if objid in context:
579            return _recursion(object), False, True
580        context[objid] = 1
581        readable = True
582        recursive = False
583        components = []
584        append = components.append
585        level += 1
586        for o in object:
587            orepr, oreadable, orecur = _safe_repr(o, context, maxlevels, level, sort_dicts)
588            append(orepr)
589            if not oreadable:
590                readable = False
591            if orecur:
592                recursive = True
593        del context[objid]
594        return format % ", ".join(components), readable, recursive
595
596    rep = repr(object)
597    return rep, (rep and not rep.startswith('<')), False
598
599_builtin_scalars = frozenset({str, bytes, bytearray, int, float, complex,
600                              bool, type(None)})
601
602def _recursion(object):
603    return ("<Recursion on %s with id=%s>"
604            % (type(object).__name__, id(object)))
605
606
607def _perfcheck(object=None):
608    import time
609    if object is None:
610        object = [("string", (1, 2), [3, 4], {5: 6, 7: 8})] * 100000
611    p = PrettyPrinter()
612    t1 = time.perf_counter()
613    _safe_repr(object, {}, None, 0, True)
614    t2 = time.perf_counter()
615    p.pformat(object)
616    t3 = time.perf_counter()
617    print("_safe_repr:", t2 - t1)
618    print("pformat:", t3 - t2)
619
620def _wrap_bytes_repr(object, width, allowance):
621    current = b''
622    last = len(object) // 4 * 4
623    for i in range(0, len(object), 4):
624        part = object[i: i+4]
625        candidate = current + part
626        if i == last:
627            width -= allowance
628        if len(repr(candidate)) > width:
629            if current:
630                yield repr(current)
631            current = part
632        else:
633            current = candidate
634    if current:
635        yield repr(current)
636
637if __name__ == "__main__":
638    _perfcheck()
639