• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1"Usage: unparse.py <path to source file>"
2import sys
3import ast
4import cStringIO
5import os
6
7# Large float and imaginary literals get turned into infinities in the AST.
8# We unparse those infinities to INFSTR.
9INFSTR = "1e" + repr(sys.float_info.max_10_exp + 1)
10
11def interleave(inter, f, seq):
12    """Call f on each item in seq, calling inter() in between.
13    """
14    seq = iter(seq)
15    try:
16        f(next(seq))
17    except StopIteration:
18        pass
19    else:
20        for x in seq:
21            inter()
22            f(x)
23
24class Unparser:
25    """Methods in this class recursively traverse an AST and
26    output source code for the abstract syntax; original formatting
27    is disregarded. """
28
29    def __init__(self, tree, file = sys.stdout):
30        """Unparser(tree, file=sys.stdout) -> None.
31         Print the source for tree to file."""
32        self.f = file
33        self.future_imports = []
34        self._indent = 0
35        self.dispatch(tree)
36        self.f.write("")
37        self.f.flush()
38
39    def fill(self, text = ""):
40        "Indent a piece of text, according to the current indentation level"
41        self.f.write("\n"+"    "*self._indent + text)
42
43    def write(self, text):
44        "Append a piece of text to the current line."
45        self.f.write(text)
46
47    def enter(self):
48        "Print ':', and increase the indentation."
49        self.write(":")
50        self._indent += 1
51
52    def leave(self):
53        "Decrease the indentation level."
54        self._indent -= 1
55
56    def dispatch(self, tree):
57        "Dispatcher function, dispatching tree type T to method _T."
58        if isinstance(tree, list):
59            for t in tree:
60                self.dispatch(t)
61            return
62        meth = getattr(self, "_"+tree.__class__.__name__)
63        meth(tree)
64
65
66    ############### Unparsing methods ######################
67    # There should be one method per concrete grammar type #
68    # Constructors should be grouped by sum type. Ideally, #
69    # this would follow the order in the grammar, but      #
70    # currently doesn't.                                   #
71    ########################################################
72
73    def _Module(self, tree):
74        for stmt in tree.body:
75            self.dispatch(stmt)
76
77    # stmt
78    def _Expr(self, tree):
79        self.fill()
80        self.dispatch(tree.value)
81
82    def _Import(self, t):
83        self.fill("import ")
84        interleave(lambda: self.write(", "), self.dispatch, t.names)
85
86    def _ImportFrom(self, t):
87        # A from __future__ import may affect unparsing, so record it.
88        if t.module and t.module == '__future__':
89            self.future_imports.extend(n.name for n in t.names)
90
91        self.fill("from ")
92        self.write("." * t.level)
93        if t.module:
94            self.write(t.module)
95        self.write(" import ")
96        interleave(lambda: self.write(", "), self.dispatch, t.names)
97
98    def _Assign(self, t):
99        self.fill()
100        for target in t.targets:
101            self.dispatch(target)
102            self.write(" = ")
103        self.dispatch(t.value)
104
105    def _AugAssign(self, t):
106        self.fill()
107        self.dispatch(t.target)
108        self.write(" "+self.binop[t.op.__class__.__name__]+"= ")
109        self.dispatch(t.value)
110
111    def _Return(self, t):
112        self.fill("return")
113        if t.value:
114            self.write(" ")
115            self.dispatch(t.value)
116
117    def _Pass(self, t):
118        self.fill("pass")
119
120    def _Break(self, t):
121        self.fill("break")
122
123    def _Continue(self, t):
124        self.fill("continue")
125
126    def _Delete(self, t):
127        self.fill("del ")
128        interleave(lambda: self.write(", "), self.dispatch, t.targets)
129
130    def _Assert(self, t):
131        self.fill("assert ")
132        self.dispatch(t.test)
133        if t.msg:
134            self.write(", ")
135            self.dispatch(t.msg)
136
137    def _Exec(self, t):
138        self.fill("exec ")
139        self.dispatch(t.body)
140        if t.globals:
141            self.write(" in ")
142            self.dispatch(t.globals)
143        if t.locals:
144            self.write(", ")
145            self.dispatch(t.locals)
146
147    def _Print(self, t):
148        self.fill("print ")
149        do_comma = False
150        if t.dest:
151            self.write(">>")
152            self.dispatch(t.dest)
153            do_comma = True
154        for e in t.values:
155            if do_comma:self.write(", ")
156            else:do_comma=True
157            self.dispatch(e)
158        if not t.nl:
159            self.write(",")
160
161    def _Global(self, t):
162        self.fill("global ")
163        interleave(lambda: self.write(", "), self.write, t.names)
164
165    def _Yield(self, t):
166        self.write("(")
167        self.write("yield")
168        if t.value:
169            self.write(" ")
170            self.dispatch(t.value)
171        self.write(")")
172
173    def _Raise(self, t):
174        self.fill('raise ')
175        if t.type:
176            self.dispatch(t.type)
177        if t.inst:
178            self.write(", ")
179            self.dispatch(t.inst)
180        if t.tback:
181            self.write(", ")
182            self.dispatch(t.tback)
183
184    def _TryExcept(self, t):
185        self.fill("try")
186        self.enter()
187        self.dispatch(t.body)
188        self.leave()
189
190        for ex in t.handlers:
191            self.dispatch(ex)
192        if t.orelse:
193            self.fill("else")
194            self.enter()
195            self.dispatch(t.orelse)
196            self.leave()
197
198    def _TryFinally(self, t):
199        if len(t.body) == 1 and isinstance(t.body[0], ast.TryExcept):
200            # try-except-finally
201            self.dispatch(t.body)
202        else:
203            self.fill("try")
204            self.enter()
205            self.dispatch(t.body)
206            self.leave()
207
208        self.fill("finally")
209        self.enter()
210        self.dispatch(t.finalbody)
211        self.leave()
212
213    def _ExceptHandler(self, t):
214        self.fill("except")
215        if t.type:
216            self.write(" ")
217            self.dispatch(t.type)
218        if t.name:
219            self.write(" as ")
220            self.dispatch(t.name)
221        self.enter()
222        self.dispatch(t.body)
223        self.leave()
224
225    def _ClassDef(self, t):
226        self.write("\n")
227        for deco in t.decorator_list:
228            self.fill("@")
229            self.dispatch(deco)
230        self.fill("class "+t.name)
231        if t.bases:
232            self.write("(")
233            for a in t.bases:
234                self.dispatch(a)
235                self.write(", ")
236            self.write(")")
237        self.enter()
238        self.dispatch(t.body)
239        self.leave()
240
241    def _FunctionDef(self, t):
242        self.write("\n")
243        for deco in t.decorator_list:
244            self.fill("@")
245            self.dispatch(deco)
246        self.fill("def "+t.name + "(")
247        self.dispatch(t.args)
248        self.write(")")
249        self.enter()
250        self.dispatch(t.body)
251        self.leave()
252
253    def _For(self, t):
254        self.fill("for ")
255        self.dispatch(t.target)
256        self.write(" in ")
257        self.dispatch(t.iter)
258        self.enter()
259        self.dispatch(t.body)
260        self.leave()
261        if t.orelse:
262            self.fill("else")
263            self.enter()
264            self.dispatch(t.orelse)
265            self.leave()
266
267    def _If(self, t):
268        self.fill("if ")
269        self.dispatch(t.test)
270        self.enter()
271        self.dispatch(t.body)
272        self.leave()
273        # collapse nested ifs into equivalent elifs.
274        while (t.orelse and len(t.orelse) == 1 and
275               isinstance(t.orelse[0], ast.If)):
276            t = t.orelse[0]
277            self.fill("elif ")
278            self.dispatch(t.test)
279            self.enter()
280            self.dispatch(t.body)
281            self.leave()
282        # final else
283        if t.orelse:
284            self.fill("else")
285            self.enter()
286            self.dispatch(t.orelse)
287            self.leave()
288
289    def _While(self, t):
290        self.fill("while ")
291        self.dispatch(t.test)
292        self.enter()
293        self.dispatch(t.body)
294        self.leave()
295        if t.orelse:
296            self.fill("else")
297            self.enter()
298            self.dispatch(t.orelse)
299            self.leave()
300
301    def _With(self, t):
302        self.fill("with ")
303        self.dispatch(t.context_expr)
304        if t.optional_vars:
305            self.write(" as ")
306            self.dispatch(t.optional_vars)
307        self.enter()
308        self.dispatch(t.body)
309        self.leave()
310
311    # expr
312    def _Str(self, tree):
313        # if from __future__ import unicode_literals is in effect,
314        # then we want to output string literals using a 'b' prefix
315        # and unicode literals with no prefix.
316        if "unicode_literals" not in self.future_imports:
317            self.write(repr(tree.s))
318        elif isinstance(tree.s, str):
319            self.write("b" + repr(tree.s))
320        elif isinstance(tree.s, unicode):
321            self.write(repr(tree.s).lstrip("u"))
322        else:
323            assert False, "shouldn't get here"
324
325    def _Name(self, t):
326        self.write(t.id)
327
328    def _Repr(self, t):
329        self.write("`")
330        self.dispatch(t.value)
331        self.write("`")
332
333    def _Num(self, t):
334        repr_n = repr(t.n)
335        # Parenthesize negative numbers, to avoid turning (-1)**2 into -1**2.
336        if repr_n.startswith("-"):
337            self.write("(")
338        # Substitute overflowing decimal literal for AST infinities.
339        self.write(repr_n.replace("inf", INFSTR))
340        if repr_n.startswith("-"):
341            self.write(")")
342
343    def _List(self, t):
344        self.write("[")
345        interleave(lambda: self.write(", "), self.dispatch, t.elts)
346        self.write("]")
347
348    def _ListComp(self, t):
349        self.write("[")
350        self.dispatch(t.elt)
351        for gen in t.generators:
352            self.dispatch(gen)
353        self.write("]")
354
355    def _GeneratorExp(self, t):
356        self.write("(")
357        self.dispatch(t.elt)
358        for gen in t.generators:
359            self.dispatch(gen)
360        self.write(")")
361
362    def _SetComp(self, t):
363        self.write("{")
364        self.dispatch(t.elt)
365        for gen in t.generators:
366            self.dispatch(gen)
367        self.write("}")
368
369    def _DictComp(self, t):
370        self.write("{")
371        self.dispatch(t.key)
372        self.write(": ")
373        self.dispatch(t.value)
374        for gen in t.generators:
375            self.dispatch(gen)
376        self.write("}")
377
378    def _comprehension(self, t):
379        self.write(" for ")
380        self.dispatch(t.target)
381        self.write(" in ")
382        self.dispatch(t.iter)
383        for if_clause in t.ifs:
384            self.write(" if ")
385            self.dispatch(if_clause)
386
387    def _IfExp(self, t):
388        self.write("(")
389        self.dispatch(t.body)
390        self.write(" if ")
391        self.dispatch(t.test)
392        self.write(" else ")
393        self.dispatch(t.orelse)
394        self.write(")")
395
396    def _Set(self, t):
397        assert(t.elts) # should be at least one element
398        self.write("{")
399        interleave(lambda: self.write(", "), self.dispatch, t.elts)
400        self.write("}")
401
402    def _Dict(self, t):
403        self.write("{")
404        def write_pair(pair):
405            (k, v) = pair
406            self.dispatch(k)
407            self.write(": ")
408            self.dispatch(v)
409        interleave(lambda: self.write(", "), write_pair, zip(t.keys, t.values))
410        self.write("}")
411
412    def _Tuple(self, t):
413        self.write("(")
414        if len(t.elts) == 1:
415            (elt,) = t.elts
416            self.dispatch(elt)
417            self.write(",")
418        else:
419            interleave(lambda: self.write(", "), self.dispatch, t.elts)
420        self.write(")")
421
422    unop = {"Invert":"~", "Not": "not", "UAdd":"+", "USub":"-"}
423    def _UnaryOp(self, t):
424        self.write("(")
425        self.write(self.unop[t.op.__class__.__name__])
426        self.write(" ")
427        # If we're applying unary minus to a number, parenthesize the number.
428        # This is necessary: -2147483648 is different from -(2147483648) on
429        # a 32-bit machine (the first is an int, the second a long), and
430        # -7j is different from -(7j).  (The first has real part 0.0, the second
431        # has real part -0.0.)
432        if isinstance(t.op, ast.USub) and isinstance(t.operand, ast.Num):
433            self.write("(")
434            self.dispatch(t.operand)
435            self.write(")")
436        else:
437            self.dispatch(t.operand)
438        self.write(")")
439
440    binop = { "Add":"+", "Sub":"-", "Mult":"*", "Div":"/", "Mod":"%",
441                    "LShift":"<<", "RShift":">>", "BitOr":"|", "BitXor":"^", "BitAnd":"&",
442                    "FloorDiv":"//", "Pow": "**"}
443    def _BinOp(self, t):
444        self.write("(")
445        self.dispatch(t.left)
446        self.write(" " + self.binop[t.op.__class__.__name__] + " ")
447        self.dispatch(t.right)
448        self.write(")")
449
450    cmpops = {"Eq":"==", "NotEq":"!=", "Lt":"<", "LtE":"<=", "Gt":">", "GtE":">=",
451                        "Is":"is", "IsNot":"is not", "In":"in", "NotIn":"not in"}
452    def _Compare(self, t):
453        self.write("(")
454        self.dispatch(t.left)
455        for o, e in zip(t.ops, t.comparators):
456            self.write(" " + self.cmpops[o.__class__.__name__] + " ")
457            self.dispatch(e)
458        self.write(")")
459
460    boolops = {ast.And: 'and', ast.Or: 'or'}
461    def _BoolOp(self, t):
462        self.write("(")
463        s = " %s " % self.boolops[t.op.__class__]
464        interleave(lambda: self.write(s), self.dispatch, t.values)
465        self.write(")")
466
467    def _Attribute(self,t):
468        self.dispatch(t.value)
469        # Special case: 3.__abs__() is a syntax error, so if t.value
470        # is an integer literal then we need to either parenthesize
471        # it or add an extra space to get 3 .__abs__().
472        if isinstance(t.value, ast.Num) and isinstance(t.value.n, int):
473            self.write(" ")
474        self.write(".")
475        self.write(t.attr)
476
477    def _Call(self, t):
478        self.dispatch(t.func)
479        self.write("(")
480        comma = False
481        for e in t.args:
482            if comma: self.write(", ")
483            else: comma = True
484            self.dispatch(e)
485        for e in t.keywords:
486            if comma: self.write(", ")
487            else: comma = True
488            self.dispatch(e)
489        if t.starargs:
490            if comma: self.write(", ")
491            else: comma = True
492            self.write("*")
493            self.dispatch(t.starargs)
494        if t.kwargs:
495            if comma: self.write(", ")
496            else: comma = True
497            self.write("**")
498            self.dispatch(t.kwargs)
499        self.write(")")
500
501    def _Subscript(self, t):
502        self.dispatch(t.value)
503        self.write("[")
504        self.dispatch(t.slice)
505        self.write("]")
506
507    # slice
508    def _Ellipsis(self, t):
509        self.write("...")
510
511    def _Index(self, t):
512        self.dispatch(t.value)
513
514    def _Slice(self, t):
515        if t.lower:
516            self.dispatch(t.lower)
517        self.write(":")
518        if t.upper:
519            self.dispatch(t.upper)
520        if t.step:
521            self.write(":")
522            self.dispatch(t.step)
523
524    def _ExtSlice(self, t):
525        interleave(lambda: self.write(', '), self.dispatch, t.dims)
526
527    # others
528    def _arguments(self, t):
529        first = True
530        # normal arguments
531        defaults = [None] * (len(t.args) - len(t.defaults)) + t.defaults
532        for a,d in zip(t.args, defaults):
533            if first:first = False
534            else: self.write(", ")
535            self.dispatch(a),
536            if d:
537                self.write("=")
538                self.dispatch(d)
539
540        # varargs
541        if t.vararg:
542            if first:first = False
543            else: self.write(", ")
544            self.write("*")
545            self.write(t.vararg)
546
547        # kwargs
548        if t.kwarg:
549            if first:first = False
550            else: self.write(", ")
551            self.write("**"+t.kwarg)
552
553    def _keyword(self, t):
554        self.write(t.arg)
555        self.write("=")
556        self.dispatch(t.value)
557
558    def _Lambda(self, t):
559        self.write("(")
560        self.write("lambda ")
561        self.dispatch(t.args)
562        self.write(": ")
563        self.dispatch(t.body)
564        self.write(")")
565
566    def _alias(self, t):
567        self.write(t.name)
568        if t.asname:
569            self.write(" as "+t.asname)
570
571def roundtrip(filename, output=sys.stdout):
572    with open(filename, "r") as pyfile:
573        source = pyfile.read()
574    tree = compile(source, filename, "exec", ast.PyCF_ONLY_AST)
575    Unparser(tree, output)
576
577
578
579def testdir(a):
580    try:
581        names = [n for n in os.listdir(a) if n.endswith('.py')]
582    except OSError:
583        sys.stderr.write("Directory not readable: %s" % a)
584    else:
585        for n in names:
586            fullname = os.path.join(a, n)
587            if os.path.isfile(fullname):
588                output = cStringIO.StringIO()
589                print 'Testing %s' % fullname
590                try:
591                    roundtrip(fullname, output)
592                except Exception as e:
593                    print '  Failed to compile, exception is %s' % repr(e)
594            elif os.path.isdir(fullname):
595                testdir(fullname)
596
597def main(args):
598    if args[0] == '--testdir':
599        for a in args[1:]:
600            testdir(a)
601    else:
602        for a in args:
603            roundtrip(a)
604
605if __name__=='__main__':
606    main(sys.argv[1:])
607