unparse.py 19.5 KB
Newer Older
Tim Peters's avatar
Tim Peters committed
1 2
"Usage: unparse.py <path to source file>"
import sys
3
import ast
4
import tokenize
5
import io
6
import os
Tim Peters's avatar
Tim Peters committed
7

8 9 10 11
# Large float and imaginary literals get turned into infinities in the AST.
# We unparse those infinities to INFSTR.
INFSTR = "1e" + repr(sys.float_info.max_10_exp + 1)

12 13 14 15 16
def interleave(inter, f, seq):
    """Call f on each item in seq, calling inter() in between.
    """
    seq = iter(seq)
    try:
17
        f(next(seq))
18 19 20 21 22 23 24
    except StopIteration:
        pass
    else:
        for x in seq:
            inter()
            f(x)

Tim Peters's avatar
Tim Peters committed
25 26 27
class Unparser:
    """Methods in this class recursively traverse an AST and
    output source code for the abstract syntax; original formatting
28
    is disregarded. """
Tim Peters's avatar
Tim Peters committed
29 30 31 32 33 34 35

    def __init__(self, tree, file = sys.stdout):
        """Unparser(tree, file=sys.stdout) -> None.
         Print the source for tree to file."""
        self.f = file
        self._indent = 0
        self.dispatch(tree)
36
        print("", file=self.f)
Tim Peters's avatar
Tim Peters committed
37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81
        self.f.flush()

    def fill(self, text = ""):
        "Indent a piece of text, according to the current indentation level"
        self.f.write("\n"+"    "*self._indent + text)

    def write(self, text):
        "Append a piece of text to the current line."
        self.f.write(text)

    def enter(self):
        "Print ':', and increase the indentation."
        self.write(":")
        self._indent += 1

    def leave(self):
        "Decrease the indentation level."
        self._indent -= 1

    def dispatch(self, tree):
        "Dispatcher function, dispatching tree type T to method _T."
        if isinstance(tree, list):
            for t in tree:
                self.dispatch(t)
            return
        meth = getattr(self, "_"+tree.__class__.__name__)
        meth(tree)


    ############### Unparsing methods ######################
    # There should be one method per concrete grammar type #
    # Constructors should be grouped by sum type. Ideally, #
    # this would follow the order in the grammar, but      #
    # currently doesn't.                                   #
    ########################################################

    def _Module(self, tree):
        for stmt in tree.body:
            self.dispatch(stmt)

    # stmt
    def _Expr(self, tree):
        self.fill()
        self.dispatch(tree.value)

82 83 84 85 86 87 88
    def _NamedExpr(self, tree):
        self.write("(")
        self.dispatch(tree.target)
        self.write(" := ")
        self.dispatch(tree.value)
        self.write(")")

Tim Peters's avatar
Tim Peters committed
89 90
    def _Import(self, t):
        self.fill("import ")
91
        interleave(lambda: self.write(", "), self.dispatch, t.names)
Tim Peters's avatar
Tim Peters committed
92

93 94
    def _ImportFrom(self, t):
        self.fill("from ")
95 96 97
        self.write("." * t.level)
        if t.module:
            self.write(t.module)
98
        self.write(" import ")
99
        interleave(lambda: self.write(", "), self.dispatch, t.names)
100

Tim Peters's avatar
Tim Peters committed
101 102 103 104 105 106 107
    def _Assign(self, t):
        self.fill()
        for target in t.targets:
            self.dispatch(target)
            self.write(" = ")
        self.dispatch(t.value)

108 109 110 111 112 113
    def _AugAssign(self, t):
        self.fill()
        self.dispatch(t.target)
        self.write(" "+self.binop[t.op.__class__.__name__]+"= ")
        self.dispatch(t.value)

114 115 116 117 118 119 120 121 122 123 124 125 126
    def _AnnAssign(self, t):
        self.fill()
        if not t.simple and isinstance(t.target, ast.Name):
            self.write('(')
        self.dispatch(t.target)
        if not t.simple and isinstance(t.target, ast.Name):
            self.write(')')
        self.write(": ")
        self.dispatch(t.annotation)
        if t.value:
            self.write(" = ")
            self.dispatch(t.value)

127
    def _Return(self, t):
128
        self.fill("return")
129
        if t.value:
130
            self.write(" ")
131 132
            self.dispatch(t.value)

133 134 135 136 137 138 139 140 141 142 143
    def _Pass(self, t):
        self.fill("pass")

    def _Break(self, t):
        self.fill("break")

    def _Continue(self, t):
        self.fill("continue")

    def _Delete(self, t):
        self.fill("del ")
144
        interleave(lambda: self.write(", "), self.dispatch, t.targets)
145 146 147 148 149 150 151 152 153

    def _Assert(self, t):
        self.fill("assert ")
        self.dispatch(t.test)
        if t.msg:
            self.write(", ")
            self.dispatch(t.msg)

    def _Global(self, t):
154 155
        self.fill("global ")
        interleave(lambda: self.write(", "), self.write, t.names)
156

157 158 159 160
    def _Nonlocal(self, t):
        self.fill("nonlocal ")
        interleave(lambda: self.write(", "), self.write, t.names)

161 162 163 164 165 166 167 168
    def _Await(self, t):
        self.write("(")
        self.write("await")
        if t.value:
            self.write(" ")
            self.dispatch(t.value)
        self.write(")")

169
    def _Yield(self, t):
170 171
        self.write("(")
        self.write("yield")
172
        if t.value:
173
            self.write(" ")
174
            self.dispatch(t.value)
175
        self.write(")")
176

177 178 179 180 181 182 183 184
    def _YieldFrom(self, t):
        self.write("(")
        self.write("yield from")
        if t.value:
            self.write(" ")
            self.dispatch(t.value)
        self.write(")")

185
    def _Raise(self, t):
186 187 188 189 190 191 192 193 194
        self.fill("raise")
        if not t.exc:
            assert not t.cause
            return
        self.write(" ")
        self.dispatch(t.exc)
        if t.cause:
            self.write(" from ")
            self.dispatch(t.cause)
195

196
    def _Try(self, t):
197 198 199 200 201 202 203 204 205 206 207
        self.fill("try")
        self.enter()
        self.dispatch(t.body)
        self.leave()
        for ex in t.handlers:
            self.dispatch(ex)
        if t.orelse:
            self.fill("else")
            self.enter()
            self.dispatch(t.orelse)
            self.leave()
208 209
        if t.finalbody:
            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 219
            self.dispatch(t.type)
        if t.name:
220 221
            self.write(" as ")
            self.write(t.name)
222 223 224 225
        self.enter()
        self.dispatch(t.body)
        self.leave()

Tim Peters's avatar
Tim Peters committed
226 227
    def _ClassDef(self, t):
        self.write("\n")
228 229 230
        for deco in t.decorator_list:
            self.fill("@")
            self.dispatch(deco)
Tim Peters's avatar
Tim Peters committed
231
        self.fill("class "+t.name)
232 233 234 235 236 237 238 239 240 241 242 243
        self.write("(")
        comma = False
        for e in t.bases:
            if comma: self.write(", ")
            else: comma = True
            self.dispatch(e)
        for e in t.keywords:
            if comma: self.write(", ")
            else: comma = True
            self.dispatch(e)
        self.write(")")

Tim Peters's avatar
Tim Peters committed
244 245 246 247 248
        self.enter()
        self.dispatch(t.body)
        self.leave()

    def _FunctionDef(self, t):
249 250 251 252 253 254
        self.__FunctionDef_helper(t, "def")

    def _AsyncFunctionDef(self, t):
        self.__FunctionDef_helper(t, "async def")

    def __FunctionDef_helper(self, t, fill_suffix):
Tim Peters's avatar
Tim Peters committed
255
        self.write("\n")
256
        for deco in t.decorator_list:
257 258
            self.fill("@")
            self.dispatch(deco)
259 260
        def_str = fill_suffix+" "+t.name + "("
        self.fill(def_str)
Tim Peters's avatar
Tim Peters committed
261
        self.dispatch(t.args)
262
        self.write(")")
263 264 265
        if t.returns:
            self.write(" -> ")
            self.dispatch(t.returns)
266 267 268 269 270
        self.enter()
        self.dispatch(t.body)
        self.leave()

    def _For(self, t):
271 272 273 274 275 276 277
        self.__For_helper("for ", t)

    def _AsyncFor(self, t):
        self.__For_helper("async for ", t)

    def __For_helper(self, fill, t):
        self.fill(fill)
278 279 280
        self.dispatch(t.target)
        self.write(" in ")
        self.dispatch(t.iter)
Tim Peters's avatar
Tim Peters committed
281 282 283
        self.enter()
        self.dispatch(t.body)
        self.leave()
284 285 286 287
        if t.orelse:
            self.fill("else")
            self.enter()
            self.dispatch(t.orelse)
288
            self.leave()
Tim Peters's avatar
Tim Peters committed
289 290 291 292 293 294 295

    def _If(self, t):
        self.fill("if ")
        self.dispatch(t.test)
        self.enter()
        self.dispatch(t.body)
        self.leave()
296 297 298 299 300 301 302 303 304 305
        # collapse nested ifs into equivalent elifs.
        while (t.orelse and len(t.orelse) == 1 and
               isinstance(t.orelse[0], ast.If)):
            t = t.orelse[0]
            self.fill("elif ")
            self.dispatch(t.test)
            self.enter()
            self.dispatch(t.body)
            self.leave()
        # final else
Tim Peters's avatar
Tim Peters committed
306 307 308 309 310 311
        if t.orelse:
            self.fill("else")
            self.enter()
            self.dispatch(t.orelse)
            self.leave()

312 313 314
    def _While(self, t):
        self.fill("while ")
        self.dispatch(t.test)
315 316 317 318 319 320 321
        self.enter()
        self.dispatch(t.body)
        self.leave()
        if t.orelse:
            self.fill("else")
            self.enter()
            self.dispatch(t.orelse)
322
            self.leave()
323

324 325
    def _With(self, t):
        self.fill("with ")
326
        interleave(lambda: self.write(", "), self.dispatch, t.items)
327 328 329 330
        self.enter()
        self.dispatch(t.body)
        self.leave()

331 332 333 334 335 336 337
    def _AsyncWith(self, t):
        self.fill("async with ")
        interleave(lambda: self.write(", "), self.dispatch, t.items)
        self.enter()
        self.dispatch(t.body)
        self.leave()

Tim Peters's avatar
Tim Peters committed
338
    # expr
339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355
    def _JoinedStr(self, t):
        self.write("f")
        string = io.StringIO()
        self._fstring_JoinedStr(t, string.write)
        self.write(repr(string.getvalue()))

    def _FormattedValue(self, t):
        self.write("f")
        string = io.StringIO()
        self._fstring_FormattedValue(t, string.write)
        self.write(repr(string.getvalue()))

    def _fstring_JoinedStr(self, t, write):
        for value in t.values:
            meth = getattr(self, "_fstring_" + type(value).__name__)
            meth(value, write)

Victor Stinner's avatar
Victor Stinner committed
356 357 358 359 360
    def _fstring_Constant(self, t, write):
        assert isinstance(t.value, str)
        value = t.value.replace("{", "{{").replace("}", "}}")
        write(value)

361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378
    def _fstring_FormattedValue(self, t, write):
        write("{")
        expr = io.StringIO()
        Unparser(t.value, expr)
        expr = expr.getvalue().rstrip("\n")
        if expr.startswith("{"):
            write(" ")  # Separate pair of opening brackets as "{ {"
        write(expr)
        if t.conversion != -1:
            conversion = chr(t.conversion)
            assert conversion in "sra"
            write(f"!{conversion}")
        if t.format_spec:
            write(":")
            meth = getattr(self, "_fstring_" + type(t.format_spec).__name__)
            meth(t.format_spec, write)
        write("}")

Tim Peters's avatar
Tim Peters committed
379 380 381
    def _Name(self, t):
        self.write(t.id)

Victor Stinner's avatar
Victor Stinner committed
382 383
    def _write_constant(self, value):
        if isinstance(value, (float, complex)):
384
            # Substitute overflowing decimal literal for AST infinities.
Victor Stinner's avatar
Victor Stinner committed
385 386 387 388 389 390 391 392 393 394 395 396 397 398
            self.write(repr(value).replace("inf", INFSTR))
        else:
            self.write(repr(value))

    def _Constant(self, t):
        value = t.value
        if isinstance(value, tuple):
            self.write("(")
            if len(value) == 1:
                self._write_constant(value[0])
                self.write(",")
            else:
                interleave(lambda: self.write(", "), self._write_constant, value)
            self.write(")")
399 400
        elif value is ...:
            self.write("...")
Victor Stinner's avatar
Victor Stinner committed
401 402 403
        else:
            self._write_constant(t.value)

Tim Peters's avatar
Tim Peters committed
404 405
    def _List(self, t):
        self.write("[")
406
        interleave(lambda: self.write(", "), self.dispatch, t.elts)
Tim Peters's avatar
Tim Peters committed
407 408
        self.write("]")

409 410 411 412 413 414 415 416 417 418 419 420 421 422
    def _ListComp(self, t):
        self.write("[")
        self.dispatch(t.elt)
        for gen in t.generators:
            self.dispatch(gen)
        self.write("]")

    def _GeneratorExp(self, t):
        self.write("(")
        self.dispatch(t.elt)
        for gen in t.generators:
            self.dispatch(gen)
        self.write(")")

423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438
    def _SetComp(self, t):
        self.write("{")
        self.dispatch(t.elt)
        for gen in t.generators:
            self.dispatch(gen)
        self.write("}")

    def _DictComp(self, t):
        self.write("{")
        self.dispatch(t.key)
        self.write(": ")
        self.dispatch(t.value)
        for gen in t.generators:
            self.dispatch(gen)
        self.write("}")

439
    def _comprehension(self, t):
440 441 442 443
        if t.is_async:
            self.write(" async for ")
        else:
            self.write(" for ")
444 445 446 447 448 449 450 451
        self.dispatch(t.target)
        self.write(" in ")
        self.dispatch(t.iter)
        for if_clause in t.ifs:
            self.write(" if ")
            self.dispatch(if_clause)

    def _IfExp(self, t):
452
        self.write("(")
453 454 455
        self.dispatch(t.body)
        self.write(" if ")
        self.dispatch(t.test)
456 457 458
        self.write(" else ")
        self.dispatch(t.orelse)
        self.write(")")
459

460 461 462 463 464 465
    def _Set(self, t):
        assert(t.elts) # should be at least one element
        self.write("{")
        interleave(lambda: self.write(", "), self.dispatch, t.elts)
        self.write("}")

466 467
    def _Dict(self, t):
        self.write("{")
468
        def write_key_value_pair(k, v):
469
            self.dispatch(k)
470
            self.write(": ")
471
            self.dispatch(v)
472 473 474 475 476 477 478 479 480 481 482

        def write_item(item):
            k, v = item
            if k is None:
                # for dictionary unpacking operator in dicts {**{'y': 2}}
                # see PEP 448 for details
                self.write("**")
                self.dispatch(v)
            else:
                write_key_value_pair(k, v)
        interleave(lambda: self.write(", "), write_item, zip(t.keys, t.values))
483 484 485 486
        self.write("}")

    def _Tuple(self, t):
        self.write("(")
487
        if len(t.elts) == 1:
Victor Stinner's avatar
Victor Stinner committed
488
            elt = t.elts[0]
489 490 491 492
            self.dispatch(elt)
            self.write(",")
        else:
            interleave(lambda: self.write(", "), self.dispatch, t.elts)
493 494
        self.write(")")

Tim Peters's avatar
Tim Peters committed
495 496 497
    unop = {"Invert":"~", "Not": "not", "UAdd":"+", "USub":"-"}
    def _UnaryOp(self, t):
        self.write("(")
498 499
        self.write(self.unop[t.op.__class__.__name__])
        self.write(" ")
Tim Peters's avatar
Tim Peters committed
500 501 502
        self.dispatch(t.operand)
        self.write(")")

503
    binop = { "Add":"+", "Sub":"-", "Mult":"*", "MatMult":"@", "Div":"/", "Mod":"%",
504
                    "LShift":"<<", "RShift":">>", "BitOr":"|", "BitXor":"^", "BitAnd":"&",
505
                    "FloorDiv":"//", "Pow": "**"}
506 507 508
    def _BinOp(self, t):
        self.write("(")
        self.dispatch(t.left)
509
        self.write(" " + self.binop[t.op.__class__.__name__] + " ")
510 511 512 513 514 515 516 517 518
        self.dispatch(t.right)
        self.write(")")

    cmpops = {"Eq":"==", "NotEq":"!=", "Lt":"<", "LtE":"<=", "Gt":">", "GtE":">=",
                        "Is":"is", "IsNot":"is not", "In":"in", "NotIn":"not in"}
    def _Compare(self, t):
        self.write("(")
        self.dispatch(t.left)
        for o, e in zip(t.ops, t.comparators):
519
            self.write(" " + self.cmpops[o.__class__.__name__] + " ")
520
            self.dispatch(e)
521
        self.write(")")
522

523
    boolops = {ast.And: 'and', ast.Or: 'or'}
524 525
    def _BoolOp(self, t):
        self.write("(")
526 527
        s = " %s " % self.boolops[t.op.__class__]
        interleave(lambda: self.write(s), self.dispatch, t.values)
528 529
        self.write(")")

530 531
    def _Attribute(self,t):
        self.dispatch(t.value)
532 533 534
        # Special case: 3.__abs__() is a syntax error, so if t.value
        # is an integer literal then we need to either parenthesize
        # it or add an extra space to get 3 .__abs__().
535
        if isinstance(t.value, ast.Constant) and isinstance(t.value.value, int):
536
            self.write(" ")
537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559
        self.write(".")
        self.write(t.attr)

    def _Call(self, t):
        self.dispatch(t.func)
        self.write("(")
        comma = False
        for e in t.args:
            if comma: self.write(", ")
            else: comma = True
            self.dispatch(e)
        for e in t.keywords:
            if comma: self.write(", ")
            else: comma = True
            self.dispatch(e)
        self.write(")")

    def _Subscript(self, t):
        self.dispatch(t.value)
        self.write("[")
        self.dispatch(t.slice)
        self.write("]")

560 561 562 563
    def _Starred(self, t):
        self.write("*")
        self.dispatch(t.value)

564
    # slice
565 566 567
    def _Ellipsis(self, t):
        self.write("...")

568 569 570 571 572 573 574 575 576 577 578 579 580
    def _Index(self, t):
        self.dispatch(t.value)

    def _Slice(self, t):
        if t.lower:
            self.dispatch(t.lower)
        self.write(":")
        if t.upper:
            self.dispatch(t.upper)
        if t.step:
            self.write(":")
            self.dispatch(t.step)

581
    def _ExtSlice(self, t):
582
        interleave(lambda: self.write(', '), self.dispatch, t.dims)
583

584 585 586 587 588 589 590
    # argument
    def _arg(self, t):
        self.write(t.arg)
        if t.annotation:
            self.write(": ")
            self.dispatch(t.annotation)

Tim Peters's avatar
Tim Peters committed
591 592 593
    # others
    def _arguments(self, t):
        first = True
594 595 596
        # normal arguments
        defaults = [None] * (len(t.args) - len(t.defaults)) + t.defaults
        for a, d in zip(t.args, defaults):
Tim Peters's avatar
Tim Peters committed
597 598 599
            if first:first = False
            else: self.write(", ")
            self.dispatch(a)
600 601 602 603 604 605
            if d:
                self.write("=")
                self.dispatch(d)

        # varargs, or bare '*' if no varargs but keyword-only arguments present
        if t.vararg or t.kwonlyargs:
Tim Peters's avatar
Tim Peters committed
606 607
            if first:first = False
            else: self.write(", ")
608 609
            self.write("*")
            if t.vararg:
610 611
                self.write(t.vararg.arg)
                if t.vararg.annotation:
612
                    self.write(": ")
613
                    self.dispatch(t.vararg.annotation)
614 615 616 617 618 619 620 621 622 623 624 625

        # keyword-only arguments
        if t.kwonlyargs:
            for a, d in zip(t.kwonlyargs, t.kw_defaults):
                if first:first = False
                else: self.write(", ")
                self.dispatch(a),
                if d:
                    self.write("=")
                    self.dispatch(d)

        # kwargs
Tim Peters's avatar
Tim Peters committed
626 627 628
        if t.kwarg:
            if first:first = False
            else: self.write(", ")
629 630
            self.write("**"+t.kwarg.arg)
            if t.kwarg.annotation:
631
                self.write(": ")
632
                self.dispatch(t.kwarg.annotation)
633 634

    def _keyword(self, t):
635 636 637 638 639
        if t.arg is None:
            self.write("**")
        else:
            self.write(t.arg)
            self.write("=")
640 641 642
        self.dispatch(t.value)

    def _Lambda(self, t):
Mark Dickinson's avatar
Mark Dickinson committed
643
        self.write("(")
644 645 646 647
        self.write("lambda ")
        self.dispatch(t.args)
        self.write(": ")
        self.dispatch(t.body)
Mark Dickinson's avatar
Mark Dickinson committed
648
        self.write(")")
Tim Peters's avatar
Tim Peters committed
649

650 651 652 653 654
    def _alias(self, t):
        self.write(t.name)
        if t.asname:
            self.write(" as "+t.asname)

655 656 657 658 659 660
    def _withitem(self, t):
        self.dispatch(t.context_expr)
        if t.optional_vars:
            self.write(" as ")
            self.dispatch(t.optional_vars)

661
def roundtrip(filename, output=sys.stdout):
662 663 664 665
    with open(filename, "rb") as pyfile:
        encoding = tokenize.detect_encoding(pyfile.readline)[0]
    with open(filename, "r", encoding=encoding) as pyfile:
        source = pyfile.read()
Mark Dickinson's avatar
Mark Dickinson committed
666
    tree = compile(source, filename, "exec", ast.PyCF_ONLY_AST)
667 668 669 670 671 672 673 674
    Unparser(tree, output)



def testdir(a):
    try:
        names = [n for n in os.listdir(a) if n.endswith('.py')]
    except OSError:
675
        print("Directory not readable: %s" % a, file=sys.stderr)
676 677 678 679
    else:
        for n in names:
            fullname = os.path.join(a, n)
            if os.path.isfile(fullname):
680 681
                output = io.StringIO()
                print('Testing %s' % fullname)
682 683
                try:
                    roundtrip(fullname, output)
684
                except Exception as e:
685
                    print('  Failed to compile, exception is %s' % repr(e))
686 687 688 689 690 691 692 693 694 695
            elif os.path.isdir(fullname):
                testdir(fullname)

def main(args):
    if args[0] == '--testdir':
        for a in args[1:]:
            testdir(a)
    else:
        for a in args:
            roundtrip(a)
Tim Peters's avatar
Tim Peters committed
696 697

if __name__=='__main__':
698
    main(sys.argv[1:])