cheesecake.ast_pp

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