unparse.py 20 KB

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