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.

254 lines
9.1 KiB

4 years ago
  1. import sys as _sys
  2. import ast as _ast
  3. from ast import boolop, cmpop, excepthandler, expr, expr_context, operator
  4. from ast import slice, stmt, unaryop, mod, AST
  5. def _make_node(Name, Fields, Attributes, Bases):
  6. def create_node(self, *args, **kwargs):
  7. nbparam = len(args) + len(kwargs)
  8. assert nbparam in (0, len(Fields)), \
  9. "Bad argument number for {}: {}, expecting {}".\
  10. format(Name, nbparam, len(Fields))
  11. self._fields = Fields
  12. self._attributes = Attributes
  13. for argname, argval in zip(self._fields, args):
  14. setattr(self, argname, argval)
  15. for argname, argval in kwargs.items():
  16. assert argname in Fields, \
  17. "Invalid Keyword argument for {}: {}".format(Name, argname)
  18. setattr(self, argname, argval)
  19. setattr(_sys.modules[__name__],
  20. Name,
  21. type(Name,
  22. Bases,
  23. {'__init__': create_node}))
  24. _nodes = {
  25. # mod
  26. 'Module': (('body',), (), (mod,)),
  27. 'Interactive': (('body',), (), (mod,)),
  28. 'Expression': (('body',), (), (mod,)),
  29. 'Suite': (('body',), (), (mod,)),
  30. # stmt
  31. 'FunctionDef': (('name', 'args', 'body', 'decorator_list', 'returns',),
  32. ('lineno', 'col_offset',),
  33. (stmt,)),
  34. 'AsyncFunctionDef': (('name', 'args', 'body',
  35. 'decorator_list', 'returns',),
  36. ('lineno', 'col_offset',),
  37. (stmt,)),
  38. 'ClassDef': (('name', 'bases', 'keywords', 'body', 'decorator_list',),
  39. ('lineno', 'col_offset',),
  40. (stmt,)),
  41. 'Return': (('value',), ('lineno', 'col_offset',),
  42. (stmt,)),
  43. 'Delete': (('targets',), ('lineno', 'col_offset',),
  44. (stmt,)),
  45. 'Assign': (('targets', 'value',), ('lineno', 'col_offset',),
  46. (stmt,)),
  47. 'AugAssign': (('target', 'op', 'value',), ('lineno', 'col_offset',),
  48. (stmt,)),
  49. 'Print': (('dest', 'values', 'nl',), ('lineno', 'col_offset',),
  50. (stmt,)),
  51. 'For': (('target', 'iter', 'body', 'orelse',), ('lineno', 'col_offset',),
  52. (stmt,)),
  53. 'AsyncFor': (('target', 'iter', 'body', 'orelse',),
  54. ('lineno', 'col_offset',),
  55. (stmt,)),
  56. 'While': (('test', 'body', 'orelse',), ('lineno', 'col_offset',),
  57. (stmt,)),
  58. 'If': (('test', 'body', 'orelse',), ('lineno', 'col_offset',),
  59. (stmt,)),
  60. 'With': (('items', 'body',), ('lineno', 'col_offset',),
  61. (stmt,)),
  62. 'AsyncWith': (('items', 'body',), ('lineno', 'col_offset',),
  63. (stmt,)),
  64. 'Raise': (('exc', 'cause',), ('lineno', 'col_offset',),
  65. (stmt,)),
  66. 'Try': (('body', 'handlers', 'orelse', 'finalbody',),
  67. ('lineno', 'col_offset',),
  68. (stmt,)),
  69. 'Assert': (('test', 'msg',), ('lineno', 'col_offset',),
  70. (stmt,)),
  71. 'Import': (('names',), ('lineno', 'col_offset',),
  72. (stmt,)),
  73. 'ImportFrom': (('module', 'names', 'level',), ('lineno', 'col_offset',),
  74. (stmt,)),
  75. 'Exec': (('body', 'globals', 'locals',), ('lineno', 'col_offset',),
  76. (stmt,)),
  77. 'Global': (('names',), ('lineno', 'col_offset',),
  78. (stmt,)),
  79. 'Nonlocal': (('names',), ('lineno', 'col_offset',),
  80. (stmt,)),
  81. 'Expr': (('value',), ('lineno', 'col_offset',),
  82. (stmt,)),
  83. 'Pass': ((), ('lineno', 'col_offset',),
  84. (stmt,)),
  85. 'Break': ((), ('lineno', 'col_offset',),
  86. (stmt,)),
  87. 'Continue': ((), ('lineno', 'col_offset',),
  88. (stmt,)),
  89. # expr
  90. 'BoolOp': (('op', 'values',), ('lineno', 'col_offset',),
  91. (expr,)),
  92. 'BinOp': (('left', 'op', 'right',), ('lineno', 'col_offset',),
  93. (expr,)),
  94. 'UnaryOp': (('op', 'operand',), ('lineno', 'col_offset',),
  95. (expr,)),
  96. 'Lambda': (('args', 'body',), ('lineno', 'col_offset',),
  97. (expr,)),
  98. 'IfExp': (('test', 'body', 'orelse',), ('lineno', 'col_offset',),
  99. (expr,)),
  100. 'Dict': (('keys', 'values',), ('lineno', 'col_offset',),
  101. (expr,)),
  102. 'Set': (('elts',), ('lineno', 'col_offset',),
  103. (expr,)),
  104. 'ListComp': (('elt', 'generators',), ('lineno', 'col_offset',),
  105. (expr,)),
  106. 'SetComp': (('elt', 'generators',), ('lineno', 'col_offset',),
  107. (expr,)),
  108. 'DictComp': (('key', 'value', 'generators',), ('lineno', 'col_offset',),
  109. (expr,)),
  110. 'GeneratorExp': (('elt', 'generators',), ('lineno', 'col_offset',),
  111. (expr,)),
  112. 'Await': (('value',), ('lineno', 'col_offset',),
  113. (expr,)),
  114. 'Yield': (('value',), ('lineno', 'col_offset',),
  115. (expr,)),
  116. 'YieldFrom': (('value',), ('lineno', 'col_offset',),
  117. (expr,)),
  118. 'Compare': (('left', 'ops', 'comparators',), ('lineno', 'col_offset',),
  119. (expr,)),
  120. 'Call': (('func', 'args', 'keywords',), ('lineno', 'col_offset',),
  121. (expr,)),
  122. 'Repr': (('value',), ('lineno', 'col_offset',),
  123. (expr,)),
  124. 'Num': (('n',), ('lineno', 'col_offset',),
  125. (expr,)),
  126. 'Str': (('s',), ('lineno', 'col_offset',),
  127. (expr,)),
  128. 'FormattedValue': (('value', 'conversion', 'format_spec',),
  129. ('lineno', 'col_offset',), (expr,)),
  130. 'JoinedStr': (('values',), ('lineno', 'col_offset',), (expr,)),
  131. 'Bytes': (('s',), ('lineno', 'col_offset',),
  132. (expr,)),
  133. 'NameConstant': (('value',), ('lineno', 'col_offset',),
  134. (expr,)),
  135. 'Ellipsis': ((), ('lineno', 'col_offset',),
  136. (expr,)),
  137. 'Attribute': (('value', 'attr', 'ctx',), ('lineno', 'col_offset',),
  138. (expr,)),
  139. 'Subscript': (('value', 'slice', 'ctx',), ('lineno', 'col_offset',),
  140. (expr,)),
  141. 'Starred': (('value', 'ctx',), ('lineno', 'col_offset',),
  142. (expr,)),
  143. 'Name': (('id', 'ctx', 'annotation'), ('lineno', 'col_offset',),
  144. (expr,)),
  145. 'List': (('elts', 'ctx',), ('lineno', 'col_offset',),
  146. (expr,)),
  147. 'Tuple': (('elts', 'ctx',), ('lineno', 'col_offset',),
  148. (expr,)),
  149. # expr_context
  150. 'Load': ((), (), (expr_context,)),
  151. 'Store': ((), (), (expr_context,)),
  152. 'Del': ((), (), (expr_context,)),
  153. 'AugLoad': ((), (), (expr_context,)),
  154. 'AugStore': ((), (), (expr_context,)),
  155. 'Param': ((), (), (expr_context,)),
  156. # slice
  157. 'Slice': (('lower', 'upper', 'step'), (), (slice,)),
  158. 'ExtSlice': (('dims',), (), (slice,)),
  159. 'Index': (('value',), (), (slice,)),
  160. # boolop
  161. 'And': ((), (), (boolop,)),
  162. 'Or': ((), (), (boolop,)),
  163. # operator
  164. 'Add': ((), (), (operator,)),
  165. 'Sub': ((), (), (operator,)),
  166. 'Mult': ((), (), (operator,)),
  167. 'MatMult': ((), (), (operator,)),
  168. 'Div': ((), (), (operator,)),
  169. 'Mod': ((), (), (operator,)),
  170. 'Pow': ((), (), (operator,)),
  171. 'LShift': ((), (), (operator,)),
  172. 'RShift': ((), (), (operator,)),
  173. 'BitOr': ((), (), (operator,)),
  174. 'BitXor': ((), (), (operator,)),
  175. 'BitAnd': ((), (), (operator,)),
  176. 'FloorDiv': ((), (), (operator,)),
  177. # unaryop
  178. 'Invert': ((), (), (unaryop, AST,)),
  179. 'Not': ((), (), (unaryop, AST,)),
  180. 'UAdd': ((), (), (unaryop, AST,)),
  181. 'USub': ((), (), (unaryop, AST,)),
  182. # cmpop
  183. 'Eq': ((), (), (cmpop,)),
  184. 'NotEq': ((), (), (cmpop,)),
  185. 'Lt': ((), (), (cmpop,)),
  186. 'LtE': ((), (), (cmpop,)),
  187. 'Gt': ((), (), (cmpop,)),
  188. 'GtE': ((), (), (cmpop,)),
  189. 'Is': ((), (), (cmpop,)),
  190. 'IsNot': ((), (), (cmpop,)),
  191. 'In': ((), (), (cmpop,)),
  192. 'NotIn': ((), (), (cmpop,)),
  193. # comprehension
  194. 'comprehension': (('target', 'iter', 'ifs', 'is_async'), (), (AST,)),
  195. # excepthandler
  196. 'ExceptHandler': (('type', 'name', 'body'), ('lineno', 'col_offset'),
  197. (excepthandler,)),
  198. # arguments
  199. 'arguments': (('args', 'vararg', 'kwonlyargs', 'kw_defaults',
  200. 'kwarg', 'defaults'), (), (AST,)),
  201. # keyword
  202. 'keyword': (('arg', 'value'), (), (AST,)),
  203. # alias
  204. 'alias': (('name', 'asname'), (), (AST,)),
  205. # withitem
  206. 'withitem': (('context_expr', 'optional_vars'), (), (AST,)),
  207. }
  208. for name, descr in _nodes.items():
  209. _make_node(name, *descr)
  210. if _sys.version_info.major == 2:
  211. from .ast2 import ast_to_gast, gast_to_ast
  212. if _sys.version_info.major == 3:
  213. from .ast3 import ast_to_gast, gast_to_ast
  214. def parse(*args, **kwargs):
  215. return ast_to_gast(_ast.parse(*args, **kwargs))
  216. def literal_eval(node_or_string):
  217. if isinstance(node_or_string, AST):
  218. node_or_string = gast_to_ast(node_or_string)
  219. return _ast.literal_eval(node_or_string)
  220. def get_docstring(node, clean=True):
  221. if not isinstance(node, (FunctionDef, ClassDef, Module)):
  222. raise TypeError("%r can't have docstrings" % node.__class__.__name__)
  223. if node.body and isinstance(node.body[0], Expr) and \
  224. isinstance(node.body[0].value, Str):
  225. if clean:
  226. import inspect
  227. return inspect.cleandoc(node.body[0].value.s)
  228. return node.body[0].value.s