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