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.decorator_list:

   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:])