You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

261 lines
7.1 KiB

4 years ago
  1. from astn import AstToGAst, GAstToAst
  2. import ast
  3. import gast
  4. class Ast2ToGAst(AstToGAst):
  5. # stmt
  6. def visit_FunctionDef(self, node):
  7. new_node = gast.FunctionDef(
  8. self._visit(node.name),
  9. self._visit(node.args),
  10. self._visit(node.body),
  11. self._visit(node.decorator_list),
  12. None, # returns
  13. )
  14. ast.copy_location(new_node, node)
  15. return new_node
  16. def visit_ClassDef(self, node):
  17. new_node = gast.ClassDef(
  18. self._visit(node.name),
  19. self._visit(node.bases),
  20. [], # keywords
  21. self._visit(node.body),
  22. self._visit(node.decorator_list),
  23. )
  24. ast.copy_location(new_node, node)
  25. return new_node
  26. def visit_With(self, node):
  27. new_node = gast.With(
  28. [gast.withitem(
  29. self._visit(node.context_expr),
  30. self._visit(node.optional_vars)
  31. )],
  32. self._visit(node.body)
  33. )
  34. ast.copy_location(new_node, node)
  35. return new_node
  36. def visit_Raise(self, node):
  37. ntype = self._visit(node.type)
  38. ninst = self._visit(node.inst)
  39. ntback = self._visit(node.tback)
  40. what = ntype
  41. if ninst is not None:
  42. what = gast.Call(ntype, [ninst], [])
  43. ast.copy_location(what, node)
  44. if ntback is not None:
  45. attr = gast.Attribute(what, 'with_traceback', gast.Load())
  46. ast.copy_location(attr, node)
  47. what = gast.Call(
  48. attr,
  49. [ntback],
  50. []
  51. )
  52. ast.copy_location(what, node)
  53. new_node = gast.Raise(what, None)
  54. ast.copy_location(new_node, node)
  55. return new_node
  56. def visit_TryExcept(self, node):
  57. new_node = gast.Try(
  58. self._visit(node.body),
  59. self._visit(node.handlers),
  60. self._visit(node.orelse),
  61. [] # finalbody
  62. )
  63. ast.copy_location(new_node, node)
  64. return new_node
  65. def visit_TryFinally(self, node):
  66. new_node = gast.Try(
  67. self._visit(node.body),
  68. [], # handlers
  69. [], # orelse
  70. self._visit(node.finalbody)
  71. )
  72. ast.copy_location(new_node, node)
  73. return new_node
  74. # expr
  75. def visit_Name(self, node):
  76. new_node = gast.Name(
  77. self._visit(node.id),
  78. self._visit(node.ctx),
  79. None,
  80. )
  81. ast.copy_location(new_node, node)
  82. return new_node
  83. def visit_Call(self, node):
  84. if node.starargs:
  85. star = gast.Starred(self._visit(node.starargs), gast.Load())
  86. ast.copy_location(star, node)
  87. starred = [star]
  88. else:
  89. starred = []
  90. if node.kwargs:
  91. kwargs = [gast.keyword(None, self._visit(node.kwargs))]
  92. else:
  93. kwargs = []
  94. new_node = gast.Call(
  95. self._visit(node.func),
  96. self._visit(node.args) + starred,
  97. self._visit(node.keywords) + kwargs,
  98. )
  99. ast.copy_location(new_node, node)
  100. return new_node
  101. def visit_comprehension(self, node):
  102. new_node = gast.comprehension(
  103. target=self._visit(node.target),
  104. iter=self._visit(node.iter),
  105. ifs=self._visit(node.ifs),
  106. is_async=0,
  107. )
  108. return ast.copy_location(new_node, node)
  109. # arguments
  110. def visit_arguments(self, node):
  111. new_node = gast.arguments(
  112. self._visit(node.args),
  113. self._visit(node.vararg),
  114. [], # kwonlyargs
  115. [], # kw_defaults
  116. self._visit(node.kwarg),
  117. self._visit(node.defaults),
  118. )
  119. return new_node
  120. class GAstToAst2(GAstToAst):
  121. # stmt
  122. def visit_FunctionDef(self, node):
  123. new_node = ast.FunctionDef(
  124. self._visit(node.name),
  125. self._visit(node.args),
  126. self._visit(node.body),
  127. self._visit(node.decorator_list),
  128. )
  129. ast.copy_location(new_node, node)
  130. return new_node
  131. def visit_ClassDef(self, node):
  132. new_node = ast.ClassDef(
  133. self._visit(node.name),
  134. self._visit(node.bases),
  135. self._visit(node.body),
  136. self._visit(node.decorator_list),
  137. )
  138. ast.copy_location(new_node, node)
  139. return new_node
  140. def visit_With(self, node):
  141. new_node = ast.With(
  142. self._visit(node.items[0].context_expr),
  143. self._visit(node.items[0].optional_vars),
  144. self._visit(node.body)
  145. )
  146. ast.copy_location(new_node, node)
  147. return new_node
  148. def visit_Raise(self, node):
  149. if isinstance(node.exc, gast.Call) and \
  150. isinstance(node.exc.func, gast.Attribute) and \
  151. node.exc.func.attr == 'with_traceback':
  152. raised = self._visit(node.exc.func.value)
  153. traceback = self._visit(node.exc.args[0])
  154. else:
  155. raised = self._visit(node.exc)
  156. traceback = None
  157. new_node = ast.Raise(raised, None, traceback)
  158. ast.copy_location(new_node, node)
  159. return new_node
  160. def visit_Try(self, node):
  161. if node.finalbody:
  162. new_node = ast.TryFinally(
  163. self._visit(node.body),
  164. self._visit(node.finalbody)
  165. )
  166. else:
  167. new_node = ast.TryExcept(
  168. self._visit(node.body),
  169. self._visit(node.handlers),
  170. self._visit(node.orelse),
  171. )
  172. ast.copy_location(new_node, node)
  173. return new_node
  174. # expr
  175. def visit_Name(self, node):
  176. new_node = ast.Name(
  177. self._visit(node.id),
  178. self._visit(node.ctx),
  179. )
  180. ast.copy_location(new_node, node)
  181. return new_node
  182. def visit_Call(self, node):
  183. if node.args and isinstance(node.args[-1], gast.Starred):
  184. args = node.args[:-1]
  185. starargs = node.args[-1].value
  186. else:
  187. args = node.args
  188. starargs = None
  189. if node.keywords and node.keywords[-1].arg is None:
  190. keywords = node.keywords[:-1]
  191. kwargs = node.keywords[-1].value
  192. else:
  193. keywords = node.keywords
  194. kwargs = None
  195. new_node = ast.Call(
  196. self._visit(node.func),
  197. self._visit(args),
  198. self._visit(keywords),
  199. self._visit(starargs),
  200. self._visit(kwargs),
  201. )
  202. ast.copy_location(new_node, node)
  203. return new_node
  204. def visit_arg(self, node):
  205. new_node = ast.Name(node.arg, ast.Param())
  206. ast.copy_location(new_node, node)
  207. return new_node
  208. # arguments
  209. def visit_arguments(self, node):
  210. new_node = ast.arguments(
  211. self._visit(node.args),
  212. self._visit(node.vararg),
  213. self._visit(node.kwarg),
  214. self._visit(node.defaults),
  215. )
  216. return new_node
  217. def ast_to_gast(node):
  218. return Ast2ToGAst().visit(node)
  219. def gast_to_ast(node):
  220. return GAstToAst2().visit(node)