Demo/parser/unparse.py

   1 "Usage: unparse.py <path to source file>"
   2 import sys
   3 import _ast
   4 import cStringIO
   5 import os
   6
   7 def interleave(inter, f, seq):
   8     """Call f on each item in seq, calling inter() in between.
   9     """
  10     seq = iter(seq)
  11     try:
  12         f(seq.next())
  13     except StopIteration:
  14         pass
  15     else:
  16         for x in seq:
  17             inter()
  18             f(x)
  19
  20 class Unparser:
  21     """Methods in this class recursively traverse an AST and
  22     output source code for the abstract syntax; original formatting
  23     is disregarged. """
  24
  25     def __init__(self, tree, file = sys.stdout):
  26         """Unparser(tree, file=sys.stdout) -> None.
  27          Print the source for tree to file."""
  28         self.f = file
  29         self._indent = 0
  30         self.dispatch(tree)
  31         print >>self.f,""
  32         self.f.flush()
  33
  34     def fill(self, text = ""):
  35         "Indent a piece of text, according to the current indentation level"
  36         self.f.write("\n"+"    "*self._indent + text)
  37
  38     def write(self, text):
  39         "Append a piece of text to the current line."
  40         self.f.write(text)
  41
  42     def enter(self):
  43         "Print ':', and increase the indentation."
  44         self.write(":")
  45         self._indent += 1
  46
  47     def leave(self):
  48         "Decrease the indentation level."
  49         self._indent -= 1
  50
  51     def dispatch(self, tree):
  52         "Dispatcher function, dispatching tree type T to method _T."
  53         if isinstance(tree, list):
  54             for t in tree:
  55                 self.dispatch(t)
  56             return
  57         meth = getattr(self, "_"+tree.__class__.__name__)
  58         meth(tree)
  59
  60
  61     ############### Unparsing methods ######################
  62     # There should be one method per concrete grammar type #
  63     # Constructors should be grouped by sum type. Ideally, #
  64     # this would follow the order in the grammar, but      #
  65     # currently doesn't.                                   #
  66     ########################################################
  67
  68     def _Module(self, tree):
  69         for stmt in tree.body:
  70             self.dispatch(stmt)
  71
  72     # stmt
  73     def _Expr(self, tree):
  74         self.fill()
  75         self.dispatch(tree.value)
  76
  77     def _Import(self, t):
  78         self.fill("import ")
  79         interleave(lambda: self.write(", "), self.dispatch, t.names)
  80
  81     def _ImportFrom(self, t):
  82         self.fill("from ")
  83         self.write(t.module)
  84         self.write(" import ")
  85         interleave(lambda: self.write(", "), self.dispatch, t.names)
  86         # XXX(jpe) what is level for?
  87
  88     def _Assign(self, t):
  89         self.fill()
  90         for target in t.targets:
  91             self.dispatch(target)
  92             self.write(" = ")
  93         self.dispatch(t.value)
  94
  95     def _AugAssign(self, t):
  96         self.fill()
  97         self.dispatch(t.target)
  98         self.write(" "+self.binop[t.op.__class__.__name__]+"= ")
  99         self.dispatch(t.value)
 100
 101     def _Return(self, t):
 102         self.fill("return")
 103         if t.value:
 104             self.write(" ")
 105             self.dispatch(t.value)
 106
 107     def _Pass(self, t):
 108         self.fill("pass")
 109
 110     def _Break(self, t):
 111         self.fill("break")
 112
 113     def _Continue(self, t):
 114         self.fill("continue")
 115
 116     def _Delete(self, t):
 117         self.fill("del ")
 118         self.dispatch(t.targets)
 119
 120     def _Assert(self, t):
 121         self.fill("assert ")
 122         self.dispatch(t.test)
 123         if t.msg:
 124             self.write(", ")
 125             self.dispatch(t.msg)
 126
 127     def _Exec(self, t):
 128         self.fill("exec ")
 129         self.dispatch(t.body)
 130         if t.globals:
 131             self.write(" in ")
 132             self.dispatch(t.globals)
 133         if t.locals:
 134             self.write(", ")
 135             self.dispatch(t.locals)
 136
 137     def _Print(self, t):
 138         self.fill("print ")
 139         do_comma = False
 140         if t.dest:
 141             self.write(">>")
 142             self.dispatch(t.dest)
 143             do_comma = True
 144         for e in t.values:
 145             if do_comma:self.write(", ")
 146             else:do_comma=True
 147             self.dispatch(e)
 148         if not t.nl:
 149             self.write(",")
 150
 151     def _Global(self, t):
 152         self.fill("global ")
 153         interleave(lambda: self.write(", "), self.write, t.names)
 154
 155     def _Yield(self, t):
 156         self.write("(")
 157         self.write("yield")
 158         if t.value:
 159             self.write(" ")
 160             self.dispatch(t.value)
 161         self.write(")")
 162
 163     def _Raise(self, t):
 164         self.fill('raise ')
 165         if t.type:
 166             self.dispatch(t.type)
 167         if t.inst:
 168             self.write(", ")
 169             self.dispatch(t.inst)
 170         if t.tback:
 171             self.write(", ")
 172             self.dispatch(t.tback)
 173
 174     def _TryExcept(self, t):
 175         self.fill("try")
 176         self.enter()
 177         self.dispatch(t.body)
 178         self.leave()
 179
 180         for ex in t.handlers:
 181             self.dispatch(ex)
 182         if t.orelse:
 183             self.fill("else")
 184             self.enter()
 185             self.dispatch(t.orelse)
 186             self.leave()
 187
 188     def _TryFinally(self, t):
 189         self.fill("try")
 190         self.enter()
 191         self.dispatch(t.body)
 192         self.leave()
 193
 194         self.fill("finally")
 195         self.enter()
 196         self.dispatch(t.finalbody)
 197         self.leave()
 198
 199     def _excepthandler(self, t):
 200         self.fill("except")
 201         if t.type:
 202             self.write(" ")
 203             self.dispatch(t.type)
 204         if t.name:
 205             self.write(", ")
 206             self.dispatch(t.name)
 207         self.enter()
 208         self.dispatch(t.body)
 209         self.leave()
 210
 211     def _ClassDef(self, t):
 212         self.write("\n")
 213         self.fill("class "+t.name)
 214         if t.bases:
 215             self.write("(")
 216             for a in t.bases:
 217                 self.dispatch(a)
 218                 self.write(", ")
 219             self.write(")")
 220         self.enter()
 221         self.dispatch(t.body)
 222         self.leave()
 223
 224     def _FunctionDef(self, t):
 225         self.write("\n")
 226         for deco in t.decorators:
 227             self.fill("@")
 228             self.dispatch(deco)
 229         self.fill("def "+t.name + "(")
 230         self.dispatch(t.args)
 231         self.write(")")
 232         self.enter()
 233         self.dispatch(t.body)
 234         self.leave()
 235
 236     def _For(self, t):
 237         self.fill("for ")
 238         self.dispatch(t.target)
 239         self.write(" in ")
 240         self.dispatch(t.iter)
 241         self.enter()
 242         self.dispatch(t.body)
 243         self.leave()
 244         if t.orelse:
 245             self.fill("else")
 246             self.enter()
 247             self.dispatch(t.orelse)
 248             self.leave
 249
 250     def _If(self, t):
 251         self.fill("if ")
 252         self.dispatch(t.test)
 253         self.enter()
 254         # XXX elif?
 255         self.dispatch(t.body)
 256         self.leave()
 257         if t.orelse:
 258             self.fill("else")
 259             self.enter()
 260             self.dispatch(t.orelse)
 261             self.leave()
 262
 263     def _While(self, t):
 264         self.fill("while ")
 265         self.dispatch(t.test)
 266         self.enter()
 267         self.dispatch(t.body)
 268         self.leave()
 269         if t.orelse:
 270             self.fill("else")
 271             self.enter()
 272             self.dispatch(t.orelse)
 273             self.leave
 274
 275     def _With(self, t):
 276         self.fill("with ")
 277         self.dispatch(t.context_expr)
 278         if t.optional_vars:
 279             self.write(" as ")
 280             self.dispatch(t.optional_vars)
 281         self.enter()
 282         self.dispatch(t.body)
 283         self.leave()
 284
 285     # expr
 286     def _Str(self, tree):
 287         self.write(repr(tree.s))
 288
 289     def _Name(self, t):
 290         self.write(t.id)
 291
 292     def _Repr(self, t):
 293         self.write("`")
 294         self.dispatch(t.value)
 295         self.write("`")
 296
 297     def _Num(self, t):
 298         self.write(repr(t.n))
 299
 300     def _List(self, t):
 301         self.write("[")
 302         interleave(lambda: self.write(", "), self.dispatch, t.elts)
 303         self.write("]")
 304
 305     def _ListComp(self, t):
 306         self.write("[")
 307         self.dispatch(t.elt)
 308         for gen in t.generators:
 309             self.dispatch(gen)
 310         self.write("]")
 311
 312     def _GeneratorExp(self, t):
 313         self.write("(")
 314         self.dispatch(t.elt)
 315         for gen in t.generators:
 316             self.dispatch(gen)
 317         self.write(")")
 318
 319     def _comprehension(self, t):
 320         self.write(" for ")
 321         self.dispatch(t.target)
 322         self.write(" in ")
 323         self.dispatch(t.iter)
 324         for if_clause in t.ifs:
 325             self.write(" if ")
 326             self.dispatch(if_clause)
 327
 328     def _IfExp(self, t):
 329         self.write("(")
 330         self.dispatch(t.body)
 331         self.write(" if ")
 332         self.dispatch(t.test)
 333         self.write(" else ")
 334         self.dispatch(t.orelse)
 335         self.write(")")
 336
 337     def _Dict(self, t):
 338         self.write("{")
 339         def writem((k, v)):
 340             self.dispatch(k)
 341             self.write(": ")
 342             self.dispatch(v)
 343         interleave(lambda: self.write(", "), writem, zip(t.keys, t.values))
 344         self.write("}")
 345
 346     def _Tuple(self, t):
 347         self.write("(")
 348         if len(t.elts) == 1:
 349             (elt,) = t.elts
 350             self.dispatch(elt)
 351             self.write(",")
 352         else:
 353             interleave(lambda: self.write(", "), self.dispatch, t.elts)
 354         self.write(")")
 355
 356     unop = {"Invert":"~", "Not": "not", "UAdd":"+", "USub":"-"}
 357     def _UnaryOp(self, t):
 358         self.write(self.unop[t.op.__class__.__name__])
 359         self.write("(")
 360         self.dispatch(t.operand)
 361         self.write(")")
 362
 363     binop = { "Add":"+", "Sub":"-", "Mult":"*", "Div":"/", "Mod":"%",
 364                     "LShift":">>", "RShift":"<<", "BitOr":"|", "BitXor":"^", "BitAnd":"&",
 365                     "FloorDiv":"//", "Pow": "**"}
 366     def _BinOp(self, t):
 367         self.write("(")
 368         self.dispatch(t.left)
 369         self.write(" " + self.binop[t.op.__class__.__name__] + " ")
 370         self.dispatch(t.right)
 371         self.write(")")
 372
 373     cmpops = {"Eq":"==", "NotEq":"!=", "Lt":"<", "LtE":"<=", "Gt":">", "GtE":">=",
 374                         "Is":"is", "IsNot":"is not", "In":"in", "NotIn":"not in"}
 375     def _Compare(self, t):
 376         self.write("(")
 377         self.dispatch(t.left)
 378         for o, e in zip(t.ops, t.comparators):
 379             self.write(" " + self.cmpops[o.__class__.__name__] + " ")
 380             self.dispatch(e)
 381             self.write(")")
 382
 383     boolops = {_ast.And: 'and', _ast.Or: 'or'}
 384     def _BoolOp(self, t):
 385         self.write("(")
 386         s = " %s " % self.boolops[t.op.__class__]
 387         interleave(lambda: self.write(s), self.dispatch, t.values)
 388         self.write(")")
 389
 390     def _Attribute(self,t):
 391         self.dispatch(t.value)
 392         self.write(".")
 393         self.write(t.attr)
 394
 395     def _Call(self, t):
 396         self.dispatch(t.func)
 397         self.write("(")
 398         comma = False
 399         for e in t.args:
 400             if comma: self.write(", ")
 401             else: comma = True
 402             self.dispatch(e)
 403         for e in t.keywords:
 404             if comma: self.write(", ")
 405             else: comma = True
 406             self.dispatch(e)
 407         if t.starargs:
 408             if comma: self.write(", ")
 409             else: comma = True
 410             self.write("*")
 411             self.dispatch(t.starargs)
 412         if t.kwargs:
 413             if comma: self.write(", ")
 414             else: comma = True
 415             self.write("**")
 416             self.dispatch(t.kwargs)
 417         self.write(")")
 418
 419     def _Subscript(self, t):
 420         self.dispatch(t.value)
 421         self.write("[")
 422         self.dispatch(t.slice)
 423         self.write("]")
 424
 425     # slice
 426     def _Ellipsis(self, t):
 427         self.write("...")
 428
 429     def _Index(self, t):
 430         self.dispatch(t.value)
 431
 432     def _Slice(self, t):
 433         if t.lower:
 434             self.dispatch(t.lower)
 435         self.write(":")
 436         if t.upper:
 437             self.dispatch(t.upper)
 438         if t.step:
 439             self.write(":")
 440             self.dispatch(t.step)
 441
 442     def _ExtSlice(self, t):
 443         interleave(lambda: self.write(', '), self.dispatch, t.dims)
 444
 445     # others
 446     def _arguments(self, t):
 447         first = True
 448         nonDef = len(t.args)-len(t.defaults)
 449         for a in t.args[0:nonDef]:
 450             if first:first = False
 451             else: self.write(", ")
 452             self.dispatch(a)
 453         for a,d in zip(t.args[nonDef:], t.defaults):
 454             if first:first = False
 455             else: self.write(", ")
 456             self.dispatch(a),
 457             self.write("=")
 458             self.dispatch(d)
 459         if t.vararg:
 460             if first:first = False
 461             else: self.write(", ")
 462             self.write("*"+t.vararg)
 463         if t.kwarg:
 464             if first:first = False
 465             else: self.write(", ")
 466             self.write("**"+t.kwarg)
 467
 468     def _keyword(self, t):
 469         self.write(t.arg)
 470         self.write("=")
 471         self.dispatch(t.value)
 472
 473     def _Lambda(self, t):
 474         self.write("lambda ")
 475         self.dispatch(t.args)
 476         self.write(": ")
 477         self.dispatch(t.body)
 478
 479     def _alias(self, t):
 480         self.write(t.name)
 481         if t.asname:
 482             self.write(" as "+t.asname)
 483
 484 def roundtrip(filename, output=sys.stdout):
 485     source = open(filename).read()
 486     tree = compile(source, filename, "exec", _ast.PyCF_ONLY_AST)
 487     Unparser(tree, output)
 488
 489
 490
 491 def testdir(a):
 492     try:
 493         names = [n for n in os.listdir(a) if n.endswith('.py')]
 494     except OSError:
 495         print >> sys.stderr, "Directory not readable: %s" % a
 496     else:
 497         for n in names:
 498             fullname = os.path.join(a, n)
 499             if os.path.isfile(fullname):
 500                 output = cStringIO.StringIO()
 501                 print 'Testing %s' % fullname
 502                 try:
 503                     roundtrip(fullname, output)
 504                 except Exception, e:
 505                     print '  Failed to compile, exception is %s' % repr(e)
 506             elif os.path.isdir(fullname):
 507                 testdir(fullname)
 508
 509 def main(args):
 510     if args[0] == '--testdir':
 511         for a in args[1:]:
 512             testdir(a)
 513     else:
 514         for a in args:
 515             roundtrip(a)
 516
 517 if __name__=='__main__':
 518     main(sys.argv[1:])