1"Usage: unparse.py <path to source file>"
2import sys
3import ast
4import tokenize
5import io
6import os
7
8# Large float and imaginary literals get turned into infinities in the AST.
9# We unparse those infinities to INFSTR.
10INFSTR = "1e" + repr(sys.float_info.max_10_exp + 1)
11
12def interleave(inter, f, seq):
13    """Call f on each item in seq, calling inter() in between.
14    """
15    seq = iter(seq)
16    try:
17        f(next(seq))
18    except StopIteration:
19        pass
20    else:
21        for x in seq:
22            inter()
23            f(x)
24
25class Unparser:
26    """Methods in this class recursively traverse an AST and
27    output source code for the abstract syntax; original formatting
28    is disregarded. """
29
30    def __init__(self, tree, file = sys.stdout):
31        """Unparser(tree, file=sys.stdout) -> None.
32         Print the source for tree to file."""
33        self.f = file
34        self._indent = 0
35        self.dispatch(tree)
36        print("", file=self.f)
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        self.fill("from ")
88        self.write("." * t.level)
89        if t.module:
90            self.write(t.module)
91        self.write(" import ")
92        interleave(lambda: self.write(", "), self.dispatch, t.names)
93
94    def _Assign(self, t):
95        self.fill()
96        for target in t.targets:
97            self.dispatch(target)
98            self.write(" = ")
99        self.dispatch(t.value)
100
101    def _AugAssign(self, t):
102        self.fill()
103        self.dispatch(t.target)
104        self.write(" "+self.binop[t.op.__class__.__name__]+"= ")
105        self.dispatch(t.value)
106
107    def _AnnAssign(self, t):
108        self.fill()
109        if not t.simple and isinstance(t.target, ast.Name):
110            self.write('(')
111        self.dispatch(t.target)
112        if not t.simple and isinstance(t.target, ast.Name):
113            self.write(')')
114        self.write(": ")
115        self.dispatch(t.annotation)
116        if t.value:
117            self.write(" = ")
118            self.dispatch(t.value)
119
120    def _Return(self, t):
121        self.fill("return")
122        if t.value:
123            self.write(" ")
124            self.dispatch(t.value)
125
126    def _Pass(self, t):
127        self.fill("pass")
128
129    def _Break(self, t):
130        self.fill("break")
131
132    def _Continue(self, t):
133        self.fill("continue")
134
135    def _Delete(self, t):
136        self.fill("del ")
137        interleave(lambda: self.write(", "), self.dispatch, t.targets)
138
139    def _Assert(self, t):
140        self.fill("assert ")
141        self.dispatch(t.test)
142        if t.msg:
143            self.write(", ")
144            self.dispatch(t.msg)
145
146    def _Global(self, t):
147        self.fill("global ")
148        interleave(lambda: self.write(", "), self.write, t.names)
149
150    def _Nonlocal(self, t):
151        self.fill("nonlocal ")
152        interleave(lambda: self.write(", "), self.write, t.names)
153
154    def _Await(self, t):
155        self.write("(")
156        self.write("await")
157        if t.value:
158            self.write(" ")
159            self.dispatch(t.value)
160        self.write(")")
161
162    def _Yield(self, t):
163        self.write("(")
164        self.write("yield")
165        if t.value:
166            self.write(" ")
167            self.dispatch(t.value)
168        self.write(")")
169
170    def _YieldFrom(self, t):
171        self.write("(")
172        self.write("yield from")
173        if t.value:
174            self.write(" ")
175            self.dispatch(t.value)
176        self.write(")")
177
178    def _Raise(self, t):
179        self.fill("raise")
180        if not t.exc:
181            assert not t.cause
182            return
183        self.write(" ")
184        self.dispatch(t.exc)
185        if t.cause:
186            self.write(" from ")
187            self.dispatch(t.cause)
188
189    def _Try(self, t):
190        self.fill("try")
191        self.enter()
192        self.dispatch(t.body)
193        self.leave()
194        for ex in t.handlers:
195            self.dispatch(ex)
196        if t.orelse:
197            self.fill("else")
198            self.enter()
199            self.dispatch(t.orelse)
200            self.leave()
201        if t.finalbody:
202            self.fill("finally")
203            self.enter()
204            self.dispatch(t.finalbody)
205            self.leave()
206
207    def _ExceptHandler(self, t):
208        self.fill("except")
209        if t.type:
210            self.write(" ")
211            self.dispatch(t.type)
212        if t.name:
213            self.write(" as ")
214            self.write(t.name)
215        self.enter()
216        self.dispatch(t.body)
217        self.leave()
218
219    def _ClassDef(self, t):
220        self.write("\n")
221        for deco in t.decorator_list:
222            self.fill("@")
223            self.dispatch(deco)
224        self.fill("class "+t.name)
225        self.write("(")
226        comma = False
227        for e in t.bases:
228            if comma: self.write(", ")
229            else: comma = True
230            self.dispatch(e)
231        for e in t.keywords:
232            if comma: self.write(", ")
233            else: comma = True
234            self.dispatch(e)
235        self.write(")")
236
237        self.enter()
238        self.dispatch(t.body)
239        self.leave()
240
241    def _FunctionDef(self, t):
242        self.__FunctionDef_helper(t, "def")
243
244    def _AsyncFunctionDef(self, t):
245        self.__FunctionDef_helper(t, "async def")
246
247    def __FunctionDef_helper(self, t, fill_suffix):
248        self.write("\n")
249        for deco in t.decorator_list:
250            self.fill("@")
251            self.dispatch(deco)
252        def_str = fill_suffix+" "+t.name + "("
253        self.fill(def_str)
254        self.dispatch(t.args)
255        self.write(")")
256        if t.returns:
257            self.write(" -> ")
258            self.dispatch(t.returns)
259        self.enter()
260        self.dispatch(t.body)
261        self.leave()
262
263    def _For(self, t):
264        self.__For_helper("for ", t)
265
266    def _AsyncFor(self, t):
267        self.__For_helper("async for ", t)
268
269    def __For_helper(self, fill, t):
270        self.fill(fill)
271        self.dispatch(t.target)
272        self.write(" in ")
273        self.dispatch(t.iter)
274        self.enter()
275        self.dispatch(t.body)
276        self.leave()
277        if t.orelse:
278            self.fill("else")
279            self.enter()
280            self.dispatch(t.orelse)
281            self.leave()
282
283    def _If(self, t):
284        self.fill("if ")
285        self.dispatch(t.test)
286        self.enter()
287        self.dispatch(t.body)
288        self.leave()
289        # collapse nested ifs into equivalent elifs.
290        while (t.orelse and len(t.orelse) == 1 and
291               isinstance(t.orelse[0], ast.If)):
292            t = t.orelse[0]
293            self.fill("elif ")
294            self.dispatch(t.test)
295            self.enter()
296            self.dispatch(t.body)
297            self.leave()
298        # final else
299        if t.orelse:
300            self.fill("else")
301            self.enter()
302            self.dispatch(t.orelse)
303            self.leave()
304
305    def _While(self, t):
306        self.fill("while ")
307        self.dispatch(t.test)
308        self.enter()
309        self.dispatch(t.body)
310        self.leave()
311        if t.orelse:
312            self.fill("else")
313            self.enter()
314            self.dispatch(t.orelse)
315            self.leave()
316
317    def _With(self, t):
318        self.fill("with ")
319        interleave(lambda: self.write(", "), self.dispatch, t.items)
320        self.enter()
321        self.dispatch(t.body)
322        self.leave()
323
324    def _AsyncWith(self, t):
325        self.fill("async with ")
326        interleave(lambda: self.write(", "), self.dispatch, t.items)
327        self.enter()
328        self.dispatch(t.body)
329        self.leave()
330
331    # expr
332    def _Bytes(self, t):
333        self.write(repr(t.s))
334
335    def _Str(self, tree):
336        self.write(repr(tree.s))
337
338    def _JoinedStr(self, t):
339        self.write("f")
340        string = io.StringIO()
341        self._fstring_JoinedStr(t, string.write)
342        self.write(repr(string.getvalue()))
343
344    def _FormattedValue(self, t):
345        self.write("f")
346        string = io.StringIO()
347        self._fstring_FormattedValue(t, string.write)
348        self.write(repr(string.getvalue()))
349
350    def _fstring_JoinedStr(self, t, write):
351        for value in t.values:
352            meth = getattr(self, "_fstring_" + type(value).__name__)
353            meth(value, write)
354
355    def _fstring_Str(self, t, write):
356        value = t.s.replace("{", "{{").replace("}", "}}")
357        write(value)
358
359    def _fstring_Constant(self, t, write):
360        assert isinstance(t.value, str)
361        value = t.value.replace("{", "{{").replace("}", "}}")
362        write(value)
363
364    def _fstring_FormattedValue(self, t, write):
365        write("{")
366        expr = io.StringIO()
367        Unparser(t.value, expr)
368        expr = expr.getvalue().rstrip("\n")
369        if expr.startswith("{"):
370            write(" ")  # Separate pair of opening brackets as "{ {"
371        write(expr)
372        if t.conversion != -1:
373            conversion = chr(t.conversion)
374            assert conversion in "sra"
375            write(f"!{conversion}")
376        if t.format_spec:
377            write(":")
378            meth = getattr(self, "_fstring_" + type(t.format_spec).__name__)
379            meth(t.format_spec, write)
380        write("}")
381
382    def _Name(self, t):
383        self.write(t.id)
384
385    def _write_constant(self, value):
386        if isinstance(value, (float, complex)):
387            self.write(repr(value).replace("inf", INFSTR))
388        else:
389            self.write(repr(value))
390
391    def _Constant(self, t):
392        value = t.value
393        if isinstance(value, tuple):
394            self.write("(")
395            if len(value) == 1:
396                self._write_constant(value[0])
397                self.write(",")
398            else:
399                interleave(lambda: self.write(", "), self._write_constant, value)
400            self.write(")")
401        else:
402            self._write_constant(t.value)
403
404    def _NameConstant(self, t):
405        self.write(repr(t.value))
406
407    def _Num(self, t):
408        # Substitute overflowing decimal literal for AST infinities.
409        self.write(repr(t.n).replace("inf", INFSTR))
410
411    def _List(self, t):
412        self.write("[")
413        interleave(lambda: self.write(", "), self.dispatch, t.elts)
414        self.write("]")
415
416    def _ListComp(self, t):
417        self.write("[")
418        self.dispatch(t.elt)
419        for gen in t.generators:
420            self.dispatch(gen)
421        self.write("]")
422
423    def _GeneratorExp(self, t):
424        self.write("(")
425        self.dispatch(t.elt)
426        for gen in t.generators:
427            self.dispatch(gen)
428        self.write(")")
429
430    def _SetComp(self, t):
431        self.write("{")
432        self.dispatch(t.elt)
433        for gen in t.generators:
434            self.dispatch(gen)
435        self.write("}")
436
437    def _DictComp(self, t):
438        self.write("{")
439        self.dispatch(t.key)
440        self.write(": ")
441        self.dispatch(t.value)
442        for gen in t.generators:
443            self.dispatch(gen)
444        self.write("}")
445
446    def _comprehension(self, t):
447        if t.is_async:
448            self.write(" async for ")
449        else:
450            self.write(" for ")
451        self.dispatch(t.target)
452        self.write(" in ")
453        self.dispatch(t.iter)
454        for if_clause in t.ifs:
455            self.write(" if ")
456            self.dispatch(if_clause)
457
458    def _IfExp(self, t):
459        self.write("(")
460        self.dispatch(t.body)
461        self.write(" if ")
462        self.dispatch(t.test)
463        self.write(" else ")
464        self.dispatch(t.orelse)
465        self.write(")")
466
467    def _Set(self, t):
468        assert(t.elts) # should be at least one element
469        self.write("{")
470        interleave(lambda: self.write(", "), self.dispatch, t.elts)
471        self.write("}")
472
473    def _Dict(self, t):
474        self.write("{")
475        def write_key_value_pair(k, v):
476            self.dispatch(k)
477            self.write(": ")
478            self.dispatch(v)
479
480        def write_item(item):
481            k, v = item
482            if k is None:
483                # for dictionary unpacking operator in dicts {**{'y': 2}}
484                # see PEP 448 for details
485                self.write("**")
486                self.dispatch(v)
487            else:
488                write_key_value_pair(k, v)
489        interleave(lambda: self.write(", "), write_item, zip(t.keys, t.values))
490        self.write("}")
491
492    def _Tuple(self, t):
493        self.write("(")
494        if len(t.elts) == 1:
495            elt = t.elts[0]
496            self.dispatch(elt)
497            self.write(",")
498        else:
499            interleave(lambda: self.write(", "), self.dispatch, t.elts)
500        self.write(")")
501
502    unop = {"Invert":"~", "Not": "not", "UAdd":"+", "USub":"-"}
503    def _UnaryOp(self, t):
504        self.write("(")
505        self.write(self.unop[t.op.__class__.__name__])
506        self.write(" ")
507        self.dispatch(t.operand)
508        self.write(")")
509
510    binop = { "Add":"+", "Sub":"-", "Mult":"*", "MatMult":"@", "Div":"/", "Mod":"%",
511                    "LShift":"<<", "RShift":">>", "BitOr":"|", "BitXor":"^", "BitAnd":"&",
512                    "FloorDiv":"//", "Pow": "**"}
513    def _BinOp(self, t):
514        self.write("(")
515        self.dispatch(t.left)
516        self.write(" " + self.binop[t.op.__class__.__name__] + " ")
517        self.dispatch(t.right)
518        self.write(")")
519
520    cmpops = {"Eq":"==", "NotEq":"!=", "Lt":"<", "LtE":"<=", "Gt":">", "GtE":">=",
521                        "Is":"is", "IsNot":"is not", "In":"in", "NotIn":"not in"}
522    def _Compare(self, t):
523        self.write("(")
524        self.dispatch(t.left)
525        for o, e in zip(t.ops, t.comparators):
526            self.write(" " + self.cmpops[o.__class__.__name__] + " ")
527            self.dispatch(e)
528        self.write(")")
529
530    boolops = {ast.And: 'and', ast.Or: 'or'}
531    def _BoolOp(self, t):
532        self.write("(")
533        s = " %s " % self.boolops[t.op.__class__]
534        interleave(lambda: self.write(s), self.dispatch, t.values)
535        self.write(")")
536
537    def _Attribute(self,t):
538        self.dispatch(t.value)
539        # Special case: 3.__abs__() is a syntax error, so if t.value
540        # is an integer literal then we need to either parenthesize
541        # it or add an extra space to get 3 .__abs__().
542        if ((isinstance(t.value, ast.Num) and isinstance(t.value.n, int))
543           or (isinstance(t.value, ast.Constant) and isinstance(t.value.value, int))):
544            self.write(" ")
545        self.write(".")
546        self.write(t.attr)
547
548    def _Call(self, t):
549        self.dispatch(t.func)
550        self.write("(")
551        comma = False
552        for e in t.args:
553            if comma: self.write(", ")
554            else: comma = True
555            self.dispatch(e)
556        for e in t.keywords:
557            if comma: self.write(", ")
558            else: comma = True
559            self.dispatch(e)
560        self.write(")")
561
562    def _Subscript(self, t):
563        self.dispatch(t.value)
564        self.write("[")
565        self.dispatch(t.slice)
566        self.write("]")
567
568    def _Starred(self, t):
569        self.write("*")
570        self.dispatch(t.value)
571
572    # slice
573    def _Ellipsis(self, t):
574        self.write("...")
575
576    def _Index(self, t):
577        self.dispatch(t.value)
578
579    def _Slice(self, t):
580        if t.lower:
581            self.dispatch(t.lower)
582        self.write(":")
583        if t.upper:
584            self.dispatch(t.upper)
585        if t.step:
586            self.write(":")
587            self.dispatch(t.step)
588
589    def _ExtSlice(self, t):
590        interleave(lambda: self.write(', '), self.dispatch, t.dims)
591
592    # argument
593    def _arg(self, t):
594        self.write(t.arg)
595        if t.annotation:
596            self.write(": ")
597            self.dispatch(t.annotation)
598
599    # others
600    def _arguments(self, t):
601        first = True
602        # normal arguments
603        defaults = [None] * (len(t.args) - len(t.defaults)) + t.defaults
604        for a, d in zip(t.args, defaults):
605            if first:first = False
606            else: self.write(", ")
607            self.dispatch(a)
608            if d:
609                self.write("=")
610                self.dispatch(d)
611
612        # varargs, or bare '*' if no varargs but keyword-only arguments present
613        if t.vararg or t.kwonlyargs:
614            if first:first = False
615            else: self.write(", ")
616            self.write("*")
617            if t.vararg:
618                self.write(t.vararg.arg)
619                if t.vararg.annotation:
620                    self.write(": ")
621                    self.dispatch(t.vararg.annotation)
622
623        # keyword-only arguments
624        if t.kwonlyargs:
625            for a, d in zip(t.kwonlyargs, t.kw_defaults):
626                if first:first = False
627                else: self.write(", ")
628                self.dispatch(a),
629                if d:
630                    self.write("=")
631                    self.dispatch(d)
632
633        # kwargs
634        if t.kwarg:
635            if first:first = False
636            else: self.write(", ")
637            self.write("**"+t.kwarg.arg)
638            if t.kwarg.annotation:
639                self.write(": ")
640                self.dispatch(t.kwarg.annotation)
641
642    def _keyword(self, t):
643        if t.arg is None:
644            self.write("**")
645        else:
646            self.write(t.arg)
647            self.write("=")
648        self.dispatch(t.value)
649
650    def _Lambda(self, t):
651        self.write("(")
652        self.write("lambda ")
653        self.dispatch(t.args)
654        self.write(": ")
655        self.dispatch(t.body)
656        self.write(")")
657
658    def _alias(self, t):
659        self.write(t.name)
660        if t.asname:
661            self.write(" as "+t.asname)
662
663    def _withitem(self, t):
664        self.dispatch(t.context_expr)
665        if t.optional_vars:
666            self.write(" as ")
667            self.dispatch(t.optional_vars)
668
669def roundtrip(filename, output=sys.stdout):
670    with open(filename, "rb") as pyfile:
671        encoding = tokenize.detect_encoding(pyfile.readline)[0]
672    with open(filename, "r", encoding=encoding) as pyfile:
673        source = pyfile.read()
674    tree = compile(source, filename, "exec", ast.PyCF_ONLY_AST)
675    Unparser(tree, output)
676
677
678
679def testdir(a):
680    try:
681        names = [n for n in os.listdir(a) if n.endswith('.py')]
682    except OSError:
683        print("Directory not readable: %s" % a, file=sys.stderr)
684    else:
685        for n in names:
686            fullname = os.path.join(a, n)
687            if os.path.isfile(fullname):
688                output = io.StringIO()
689                print('Testing %s' % fullname)
690                try:
691                    roundtrip(fullname, output)
692                except Exception as e:
693                    print('  Failed to compile, exception is %s' % repr(e))
694            elif os.path.isdir(fullname):
695                testdir(fullname)
696
697def main(args):
698    if args[0] == '--testdir':
699        for a in args[1:]:
700            testdir(a)
701    else:
702        for a in args:
703            roundtrip(a)
704
705if __name__=='__main__':
706    main(sys.argv[1:])
707