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