blob: 3a3fc97a0184c39678e735c488bb2ca787e16ec2 [file] [log] [blame]
Zack Williams9731cdc2019-11-22 15:42:30 -07001# Copyright 2017-present Open Networking Foundation
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# Copyright 2013, Big Switch Networks, Inc.
15
16"""
17pp - port of Ruby's PP library
18Also based on Lindig, C., & GbR, G. D. (2000). Strictly Pretty.
19
20Example usage:
21>>> import pp.pp as pp
22>>> print pp([[1, 2], [3, 4]], maxwidth=15)
23[
24 [ 1, 2 ],
25 [ 3, 4 ]
26]
27"""
28import unittest
29from contextlib import contextmanager
30
31def pp(obj, maxwidth=79):
32 """
33 Pretty-print the given object.
34 """
35 ctx = PrettyPrinter(maxwidth=maxwidth)
36 ctx.pp(obj)
37 return str(ctx)
38
39
40## Pretty-printers for builtin classes
41
42def pretty_print_list(pp, obj):
43 with pp.group():
44 pp.text('[')
45 with pp.indent(2):
46 for v in obj:
47 if not pp.first(): pp.text(',')
48 pp.breakable()
49 pp.pp(v)
50 pp.breakable()
51 pp.text(']')
52
53def pretty_print_dict(pp, obj):
54 with pp.group():
55 pp.text('{')
56 with pp.indent(2):
57 for (k, v) in sorted(obj.items()):
58 if not pp.first(): pp.text(',')
59 pp.breakable()
60 pp.pp(k)
61 pp.text(': ')
62 pp.pp(v)
63 pp.breakable()
64 pp.text('}')
65
66pretty_printers = {
67 list: pretty_print_list,
68 dict: pretty_print_dict,
69}
70
71
72## Implementation
73
74class PrettyPrinter(object):
75 def __init__(self, maxwidth):
76 self.maxwidth = maxwidth
77 self.cur_indent = 0
78 self.root_group = Group()
79 self.group_stack = [self.root_group]
80
81 def current_group(self):
82 return self.group_stack[-1]
83
84 def text(self, s):
85 self.current_group().append(str(s))
86
87 def breakable(self, sep=' '):
88 self.current_group().append(Breakable(sep, self.cur_indent))
89
90 def first(self):
91 return self.current_group().first()
92
93 @contextmanager
94 def indent(self, n):
95 self.cur_indent += n
96 yield
97 self.cur_indent -= n
98
99 @contextmanager
100 def group(self):
101 self.group_stack.append(Group())
102 yield
103 new_group = self.group_stack.pop()
104 self.current_group().append(new_group)
105
106 def pp(self, obj):
107 if hasattr(obj, "pretty_print"):
108 obj.pretty_print(self)
109 elif type(obj) in pretty_printers:
110 pretty_printers[type(obj)](self, obj)
111 else:
112 self.text(repr(obj))
113
114 def __str__(self):
115 return self.root_group.render(0, self.maxwidth)
116
117class Group(object):
118 __slots__ = ["fragments", "length", "_first"]
119
120 def __init__(self):
121 self.fragments = []
122 self.length = 0
123 self._first = True
124
125 def append(self, x):
126 self.fragments.append(x)
127 self.length += len(x)
128
129 def first(self):
130 if self._first:
131 self._first = False
132 return True
133 return False
134
135 def __len__(self):
136 return self.length
137
138 def render(self, curwidth, maxwidth):
139 dobreak = len(self) > (maxwidth - curwidth)
140
141 a = []
142 for x in self.fragments:
143 if isinstance(x, Breakable):
144 if dobreak:
145 a.append('\n')
146 a.append(' ' * x.indent)
147 curwidth = 0
148 else:
149 a.append(x.sep)
150 elif isinstance(x, Group):
151 a.append(x.render(curwidth, maxwidth))
152 else:
153 a.append(x)
154 curwidth += len(a[-1])
155 return ''.join(a)
156
157class Breakable(object):
158 __slots__ = ["sep", "indent"]
159
160 def __init__(self, sep, indent):
161 self.sep = sep
162 self.indent = indent
163
164 def __len__(self):
165 return len(self.sep)
166
167
168## Tests
169
170class TestPP(unittest.TestCase):
171 def test_scalars(self):
172 self.assertEquals(pp(1), "1")
173 self.assertEquals(pp("foo"), "'foo'")
174
175 def test_hash(self):
176 expected = """{ 1: 'a', 'b': 2 }"""
177 self.assertEquals(pp(eval(expected)), expected)
178 expected = """\
179{
180 1: 'a',
181 'b': 2
182}"""
183 self.assertEquals(pp(eval(expected), maxwidth=0), expected)
184
185 def test_array(self):
186 expected = """[ 1, 'a', 2 ]"""
187 self.assertEquals(pp(eval(expected)), expected)
188 expected = """\
189[
190 1,
191 'a',
192 2
193]"""
194 self.assertEquals(pp(eval(expected), maxwidth=0), expected)
195
196 def test_nested(self):
197 expected = """[ [ 1, 2 ], [ 3, 4 ] ]"""
198 self.assertEquals(pp(eval(expected)), expected)
199 expected = """\
200[
201 [
202 1,
203 2
204 ],
205 [
206 3,
207 4
208 ]
209]"""
210 self.assertEquals(pp(eval(expected), maxwidth=0), expected)
211
212 def test_breaking(self):
213 expected = """\
214[
215 [ 1, 2 ],
216 'abcdefghijklmnopqrstuvwxyz'
217]"""
218 self.assertEquals(pp(eval(expected), maxwidth=24), expected)
219 expected = """\
220[
221 [ 'abcd', 2 ],
222 [ '0123456789' ],
223 [
224 '0123456789',
225 'abcdefghij'
226 ],
227 [ 'abcdefghijklmnop' ],
228 [
229 'abcdefghijklmnopq'
230 ],
231 { 'k': 'v' },
232 {
233 1: [ 2, [ 3, 4 ] ],
234 'foo': 'abcdefghijklmnop'
235 }
236]"""
237 self.assertEquals(pp(eval(expected), maxwidth=24), expected)
238 expected = """\
239[
240 [ 1, 2 ],
241 [ 3, 4 ]
242]"""
243 self.assertEquals(pp(eval(expected), maxwidth=15), expected)
244
245 # This is an edge case where our simpler algorithm breaks down.
246 @unittest.expectedFailure
247 def test_greedy_breaking(self):
248 expected = """\
249abc def
250ghijklmnopqrstuvwxyz\
251"""
252 pp = PrettyPrinter(maxwidth=8)
253 pp.text("abc")
254 with pp.group():
255 pp.breakable()
256 pp.text("def")
257 with pp.group():
258 pp.breakable()
259 pp.text("ghijklmnopqrstuvwxyz")
260 self.assertEquals(str(pp), expected)
261
262if __name__ == '__main__':
263 unittest.main()