root/branches/mk/cheesecake/ast_pp.py

Revision 31, 7.8 kB (checked in by mk, 4 years ago)

Use the latest pydoctor (closes ticket #8).

  • Property svn:eol-style set to native
  • Property svn:keywords set to Author Date Id Revision
Line 
1 # this is stolen from exarkun's sandbox
2 import sys
3
4 from StringIO import StringIO
5 from pprint import pprint
6 from compiler import parse, walk
7
8 class SourceWriter(object):
9     _i = 0
10
11     def __init__(self):
12         self.s = StringIO()
13
14     def w(self, s):
15         self.s.write(s)
16
17     def nl(self):
18         self.s.write('\n')
19         self.s.write(' ' * 4 * self._i)
20
21     def indent(self):
22         self._i += 1
23         self.nl()
24
25     def dedent(self):
26         self._i -= 1
27         self.nl()
28
29     def visitModule(self, node):
30         if node.doc is not None:
31             self.wl(repr(node.doc))
32         walk(node.node, self)
33
34     def visitStmt(self, node):
35         for n in node.getChildren():
36             walk(n, self)
37
38     def visitFunction(self, node):
39         fmt = 'def %s(%s):'
40         if node.defaults:
41             nargs = len(node.argnames)
42             ndefs = len(node.defaults)
43             noDefaults = node.argnames[:nargs-ndefs]
44             s = ', '.join(node.argnames[:noDefaults])
45             if ndefs < nargs:
46                 argdefs = zip(node.argnames[noDefaults:], node.defaults)
47                 s = s + ', ' + ', '.join(['='.join(x) for x in argdefs])
48         else:
49             s = ', '.join(node.argnames)
50         self.w(fmt % (node.name, s))
51         self.indent()
52         try:
53             walk(node.code, self)
54         finally:
55             self.dedent()
56
57     def visitAssign(self, node):
58         walk(node.nodes[0], self)
59         self.w(' = ')
60         walk(node.expr, self)
61         self.nl()
62
63     def visitAssName(self, node):
64         self.w(node.name)
65
66     def visitCallFunc(self, node):
67         walk(node.node, self)
68         self.w('(')
69         for a in node.args[:-1]:
70             walk(a, self)
71             self.w(', ')
72         for a in node.args[-1:]:
73             walk(a, self)
74         self.w(')')
75
76     def visitListComp(self, node):
77         self.w('[')
78         walk(node.expr, self)
79         for q in node.quals:
80             walk(q, self)
81         self.w(']')
82
83     def visitListCompFor(self, node):
84         self.w(' for ')
85         walk(node.assign, self)
86         self.w(' in ')
87         walk(node.list, self)
88         for expr in node.ifs:
89             self.w(' if ')
90             walk(expr, self)
91
92     def visitName(self, node):
93         self.w(node.name)
94
95     def visitDiscard(self, node):
96         walk(node.expr, self)
97         self.nl()
98
99     def visitPrintnl(self, node):
100         self.w('print ')
101         if node.dest:
102             self.w('>>')
103             walk(node.dest, self)
104             self.w(', ')
105         for e in node.nodes:
106             walk(e, self)
107         self.nl()
108
109     def visitGetattr(self, node):
110         walk(node.expr, self)
111         self.w('.')
112         self.w(node.attrname)
113
114     def visitImport(self, node):
115         self.w('import ')
116         for (mod, as) in node.names:
117             self.w(mod)
118             if as is not None:
119                 self.w(' as ')
120                 self.w(as)
121             self.w(', ')
122         self.nl()
123
124     def visitFrom(self, node):
125         self.w('from ')
126         self.w(node.modname)
127         self.w(' import ')
128         for (mod, as) in node.names:
129             self.w(mod)
130             if as is not None:
131                 self.w(' as ')
132                 self.w(as)
133             self.w(', ')
134         self.nl()
135
136     def visitConst(self, node):
137         self.w(repr(node.value))
138
139     def visitReturn(self, node):
140         self.w('return ')
141         walk(node.value, self)
142         self.nl()
143
144     def visitClass(self, node):
145         self.w('class ')
146         self.w(node.name)
147         if node.bases:
148             self.w('(')
149             for b in node.bases:
150                 walk(b, self)
151                 self.w(', ')
152             self.w('):')
153         self.indent()
154         try:
155             if node.doc is not None:
156                 self.w(repr(node.doc))
157             walk(node.code, self)
158         finally:
159             self.dedent()
160
161     def visitAssAttr(self, node):
162         walk(node.expr, self)
163         self.w('.')
164         self.w(node.attrname)
165
166     def visitMul(self, node):
167         walk(node.left, self)
168         self.w(' * ')
169         walk(node.right, self)
170
171     def visitSub(self, node):
172         walk(node.left, self)
173         self.w(' - ')
174         walk(node.right, self)
175
176     def visitAdd(self, node):
177         walk(node.left, self)
178         self.w(' + ')
179         walk(node.right, self)
180
181     def visitMod(self, node):
182         walk(node.left, self)
183         self.w(' % ')
184         walk(node.right, self)
185
186     def visitAugAssign(self, node):
187         walk(node.node, self)
188         self.w(' ')
189         self.w(node.op)
190         self.w(' ')
191         walk(node.expr, self)
192         self.nl()
193
194     def visitIf(self, node):
195         keyword = 'if'
196         for (cond, body) in node.tests:
197             self.w(keyword)
198             self.w(' ')
199             walk(cond, self)
200             self.w(':')
201             self.indent()
202             try:
203                 walk(body, self)
204             finally:
205                 self.dedent()
206             keyword = 'elif'
207         if node.else_:
208             self.w('else:')
209             self.indent()
210             try:
211                 walk(node.else_, self)
212             finally:
213                 self.dedent()
214
215     def visitCompare(self, node):
216         walk(node.expr, self)
217         for (op, arg) in node.ops:
218             self.w(' ')
219             self.w(op)
220             self.w(' ')
221             walk(arg, self)
222
223     def visitFor(self, node):
224         self.w('for ')
225         walk(node.assign, self)
226         self.w(' in ')
227         walk(node.list, self)
228         self.w(':')
229         self.indent()
230         try:
231             walk(node.body, self)
232         finally:
233             self.dedent()
234         if node.else_:
235             self.w('else:')
236             self.indent()
237             try:
238                 walk(node.else_, self)
239             finally:
240                 self.dedent()
241
242     def visitSlice(self, node):
243         walk(node.expr, self)
244         self.w('[')
245         if node.lower:
246             walk(node.lower, self)
247         self.w(':')
248         if node.upper:
249             walk(node.upper, self)
250         self.w(']')
251
252     def visitTuple(self, node):
253         self.w('(')
254         if len(node.nodes) == 0:
255             pass
256         elif len(node.nodes) == 1:
257             walk(node.nodes[0], self)
258             self.w(',')
259         else:
260             for expr in node.nodes[:-1]:
261                 walk(expr, self)
262                 self.w(', ')
263             walk(node.nodes[-1], self)
264         self.w(')')
265
266     def visitTryFinally(self, node):
267         self.w('try:')
268         self.indent()
269         try:
270             walk(node.body, self)
271         finally:
272             self.dedent()
273         self.w('finally:')
274         self.indent()
275         try:
276             walk(node.final, self)
277         finally:
278             self.dedent()
279
280     def visitSubscript(self, node):
281         walk(node.expr, self)
282         self.w('[')
283         walk(node.subs[0], self)
284         self.w(']')
285
286     def visitUnarySub(self, node):
287         self.w('-')
288         walk(node.expr, self)
289
290     def visitAssTuple(self, node):
291         self.w('(')
292         for expr in node.nodes:
293             walk(expr, self)
294             self.w(', ')
295         self.w(')')
296
297     def visitRaise(self, node):
298         self.w('raise ')
299         walk(node.expr1, self)
300         if node.expr2:
301             self.w(', ')
302             walk(node.expr2, self)
303             if node.expr3:
304                 self.w(', ')
305                 walk(node.expr3, self)
306         self.nl()
307
308     def visitDict(self, node):
309         self.w('{')
310         for (k, v) in node.items:
311             walk(k, self)
312             self.w(':')
313             walk(v, self)
314             self.w(',')
315         self.w('}')
316
317     def __str__(self):
318         return self.s.getvalue()
319
320 def pp(ast):
321     sw = SourceWriter()
322     walk(ast, sw)
323     return sw.s.getvalue()
324
325 def magic(s):
326     ast = parse(s)
327     sw = SourceWriter()
328     walk(ast, sw)
329     return ast, sw
330
331 if __name__ == '__main__':
332     f = file(__file__, 'r').read()
333     ast, sw = magic(f)
334     print sw
335     print ast
Note: See TracBrowser for help on using the browser.