Package cheesecake :: Module ast_pp
[hide private]
[frames] | no frames]

Source Code for Module cheesecake.ast_pp

  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 336