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 sys as _sys
38import warnings
39
40try:
41    from cStringIO import StringIO as _StringIO
42except ImportError:
43    from StringIO import StringIO as _StringIO
44
45__all__ = ["pprint","pformat","isreadable","isrecursive","saferepr",
46           "PrettyPrinter"]
47
48# cache these for faster access:
49_commajoin = ", ".join
50_id = id
51_len = len
52_type = type
53
54
55def pprint(object, stream=None, indent=1, width=80, depth=None):
56    """Pretty-print a Python object to a stream [default is sys.stdout]."""
57    printer = PrettyPrinter(
58        stream=stream, indent=indent, width=width, depth=depth)
59    printer.pprint(object)
60
61def pformat(object, indent=1, width=80, depth=None):
62    """Format a Python object into a pretty-printed representation."""
63    return PrettyPrinter(indent=indent, width=width, depth=depth).pformat(object)
64
65def saferepr(object):
66    """Version of repr() which can handle recursive data structures."""
67    return _safe_repr(object, {}, None, 0)[0]
68
69def isreadable(object):
70    """Determine if saferepr(object) is readable by eval()."""
71    return _safe_repr(object, {}, None, 0)[1]
72
73def isrecursive(object):
74    """Determine if object requires a recursive representation."""
75    return _safe_repr(object, {}, None, 0)[2]
76
77def _sorted(iterable):
78    with warnings.catch_warnings():
79        if _sys.py3kwarning:
80            warnings.filterwarnings("ignore", "comparing unequal types "
81                                    "not supported", DeprecationWarning)
82        return sorted(iterable)
83
84class PrettyPrinter:
85    def __init__(self, indent=1, width=80, depth=None, stream=None):
86        """Handle pretty printing operations onto a stream using a set of
87        configured parameters.
88
89        indent
90            Number of spaces to indent for each level of nesting.
91
92        width
93            Attempted maximum number of columns in the output.
94
95        depth
96            The maximum depth to print out nested structures.
97
98        stream
99            The desired output stream.  If omitted (or false), the standard
100            output stream available at construction will be used.
101
102        """
103        indent = int(indent)
104        width = int(width)
105        assert indent >= 0, "indent must be >= 0"
106        assert depth is None or depth > 0, "depth must be > 0"
107        assert width, "width must be != 0"
108        self._depth = depth
109        self._indent_per_level = indent
110        self._width = width
111        if stream is not None:
112            self._stream = stream
113        else:
114            self._stream = _sys.stdout
115
116    def pprint(self, object):
117        self._format(object, self._stream, 0, 0, {}, 0)
118        self._stream.write("\n")
119
120    def pformat(self, object):
121        sio = _StringIO()
122        self._format(object, sio, 0, 0, {}, 0)
123        return sio.getvalue()
124
125    def isrecursive(self, object):
126        return self.format(object, {}, 0, 0)[2]
127
128    def isreadable(self, object):
129        s, readable, recursive = self.format(object, {}, 0, 0)
130        return readable and not recursive
131
132    def _format(self, object, stream, indent, allowance, context, level):
133        level = level + 1
134        objid = _id(object)
135        if objid in context:
136            stream.write(_recursion(object))
137            self._recursive = True
138            self._readable = False
139            return
140        rep = self._repr(object, context, level - 1)
141        typ = _type(object)
142        sepLines = _len(rep) > (self._width - 1 - indent - allowance)
143        write = stream.write
144
145        if self._depth and level > self._depth:
146            write(rep)
147            return
148
149        r = getattr(typ, "__repr__", None)
150        if issubclass(typ, dict) and r is dict.__repr__:
151            write('{')
152            if self._indent_per_level > 1:
153                write((self._indent_per_level - 1) * ' ')
154            length = _len(object)
155            if length:
156                context[objid] = 1
157                indent = indent + self._indent_per_level
158                items = _sorted(object.items())
159                key, ent = items[0]
160                rep = self._repr(key, context, level)
161                write(rep)
162                write(': ')
163                self._format(ent, stream, indent + _len(rep) + 2,
164                              allowance + 1, context, level)
165                if length > 1:
166                    for key, ent in items[1:]:
167                        rep = self._repr(key, context, level)
168                        if sepLines:
169                            write(',\n%s%s: ' % (' '*indent, rep))
170                        else:
171                            write(', %s: ' % rep)
172                        self._format(ent, stream, indent + _len(rep) + 2,
173                                      allowance + 1, context, level)
174                indent = indent - self._indent_per_level
175                del context[objid]
176            write('}')
177            return
178
179        if ((issubclass(typ, list) and r is list.__repr__) or
180            (issubclass(typ, tuple) and r is tuple.__repr__) or
181            (issubclass(typ, set) and r is set.__repr__) or
182            (issubclass(typ, frozenset) and r is frozenset.__repr__)
183           ):
184            length = _len(object)
185            if issubclass(typ, list):
186                write('[')
187                endchar = ']'
188            elif issubclass(typ, set):
189                if not length:
190                    write('set()')
191                    return
192                write('set([')
193                endchar = '])'
194                object = _sorted(object)
195                indent += 4
196            elif issubclass(typ, frozenset):
197                if not length:
198                    write('frozenset()')
199                    return
200                write('frozenset([')
201                endchar = '])'
202                object = _sorted(object)
203                indent += 10
204            else:
205                write('(')
206                endchar = ')'
207            if self._indent_per_level > 1 and sepLines:
208                write((self._indent_per_level - 1) * ' ')
209            if length:
210                context[objid] = 1
211                indent = indent + self._indent_per_level
212                self._format(object[0], stream, indent, allowance + 1,
213                             context, level)
214                if length > 1:
215                    for ent in object[1:]:
216                        if sepLines:
217                            write(',\n' + ' '*indent)
218                        else:
219                            write(', ')
220                        self._format(ent, stream, indent,
221                                      allowance + 1, context, level)
222                indent = indent - self._indent_per_level
223                del context[objid]
224            if issubclass(typ, tuple) and length == 1:
225                write(',')
226            write(endchar)
227            return
228
229        write(rep)
230
231    def _repr(self, object, context, level):
232        repr, readable, recursive = self.format(object, context.copy(),
233                                                self._depth, level)
234        if not readable:
235            self._readable = False
236        if recursive:
237            self._recursive = True
238        return repr
239
240    def format(self, object, context, maxlevels, level):
241        """Format object for a specific context, returning a string
242        and flags indicating whether the representation is 'readable'
243        and whether the object represents a recursive construct.
244        """
245        return _safe_repr(object, context, maxlevels, level)
246
247
248# Return triple (repr_string, isreadable, isrecursive).
249
250def _safe_repr(object, context, maxlevels, level):
251    typ = _type(object)
252    if typ is str:
253        if 'locale' not in _sys.modules:
254            return repr(object), True, False
255        if "'" in object and '"' not in object:
256            closure = '"'
257            quotes = {'"': '\\"'}
258        else:
259            closure = "'"
260            quotes = {"'": "\\'"}
261        qget = quotes.get
262        sio = _StringIO()
263        write = sio.write
264        for char in object:
265            if char.isalpha():
266                write(char)
267            else:
268                write(qget(char, repr(char)[1:-1]))
269        return ("%s%s%s" % (closure, sio.getvalue(), closure)), True, False
270
271    r = getattr(typ, "__repr__", None)
272    if issubclass(typ, dict) and r is dict.__repr__:
273        if not object:
274            return "{}", True, False
275        objid = _id(object)
276        if maxlevels and level >= maxlevels:
277            return "{...}", False, objid in context
278        if objid in context:
279            return _recursion(object), False, True
280        context[objid] = 1
281        readable = True
282        recursive = False
283        components = []
284        append = components.append
285        level += 1
286        saferepr = _safe_repr
287        for k, v in _sorted(object.items()):
288            krepr, kreadable, krecur = saferepr(k, context, maxlevels, level)
289            vrepr, vreadable, vrecur = saferepr(v, context, maxlevels, level)
290            append("%s: %s" % (krepr, vrepr))
291            readable = readable and kreadable and vreadable
292            if krecur or vrecur:
293                recursive = True
294        del context[objid]
295        return "{%s}" % _commajoin(components), readable, recursive
296
297    if (issubclass(typ, list) and r is list.__repr__) or \
298       (issubclass(typ, tuple) and r is tuple.__repr__):
299        if issubclass(typ, list):
300            if not object:
301                return "[]", True, False
302            format = "[%s]"
303        elif _len(object) == 1:
304            format = "(%s,)"
305        else:
306            if not object:
307                return "()", True, False
308            format = "(%s)"
309        objid = _id(object)
310        if maxlevels and level >= maxlevels:
311            return format % "...", False, objid in context
312        if objid in context:
313            return _recursion(object), False, True
314        context[objid] = 1
315        readable = True
316        recursive = False
317        components = []
318        append = components.append
319        level += 1
320        for o in object:
321            orepr, oreadable, orecur = _safe_repr(o, context, maxlevels, level)
322            append(orepr)
323            if not oreadable:
324                readable = False
325            if orecur:
326                recursive = True
327        del context[objid]
328        return format % _commajoin(components), readable, recursive
329
330    rep = repr(object)
331    return rep, (rep and not rep.startswith('<')), False
332
333
334def _recursion(object):
335    return ("<Recursion on %s with id=%s>"
336            % (_type(object).__name__, _id(object)))
337
338
339def _perfcheck(object=None):
340    import time
341    if object is None:
342        object = [("string", (1, 2), [3, 4], {5: 6, 7: 8})] * 100000
343    p = PrettyPrinter()
344    t1 = time.time()
345    _safe_repr(object, {}, None, 0)
346    t2 = time.time()
347    p.pformat(object)
348    t3 = time.time()
349    print "_safe_repr:", t2 - t1
350    print "pformat:", t3 - t2
351
352if __name__ == "__main__":
353    _perfcheck()
354