383 lines
15 KiB
Python
383 lines
15 KiB
Python
import sys as _sys
|
|
import ast as _ast
|
|
from ast import boolop, cmpop, excepthandler, expr, expr_context, operator
|
|
from ast import slice, stmt, unaryop, mod, AST
|
|
from ast import iter_child_nodes, walk
|
|
|
|
try:
|
|
from ast import TypeIgnore
|
|
except ImportError:
|
|
class TypeIgnore(AST):
|
|
pass
|
|
|
|
|
|
def _make_node(Name, Fields, Attributes, Bases):
|
|
def create_node(self, *args, **kwargs):
|
|
nbparam = len(args) + len(kwargs)
|
|
assert nbparam in (0, len(Fields)), \
|
|
"Bad argument number for {}: {}, expecting {}".\
|
|
format(Name, nbparam, len(Fields))
|
|
self._fields = Fields
|
|
self._attributes = Attributes
|
|
for argname, argval in zip(self._fields, args):
|
|
setattr(self, argname, argval)
|
|
for argname, argval in kwargs.items():
|
|
assert argname in Fields, \
|
|
"Invalid Keyword argument for {}: {}".format(Name, argname)
|
|
setattr(self, argname, argval)
|
|
|
|
setattr(_sys.modules[__name__],
|
|
Name,
|
|
type(Name,
|
|
Bases,
|
|
{'__init__': create_node}))
|
|
|
|
|
|
_nodes = (
|
|
# mod
|
|
('Module', (('body', 'type_ignores'), (), (mod,))),
|
|
('Interactive', (('body',), (), (mod,))),
|
|
('Expression', (('body',), (), (mod,))),
|
|
('FunctionType', (('argtypes', 'returns'), (), (mod,))),
|
|
('Suite', (('body',), (), (mod,))),
|
|
|
|
# stmt
|
|
('FunctionDef', (('name', 'args', 'body', 'decorator_list', 'returns',
|
|
'type_comment'),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('AsyncFunctionDef', (('name', 'args', 'body',
|
|
'decorator_list', 'returns',
|
|
'type_comment'),
|
|
('lineno', 'col_offset',
|
|
'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('ClassDef', (('name', 'bases', 'keywords', 'body', 'decorator_list',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Return', (('value',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Delete', (('targets',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Assign', (('targets', 'value',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('AugAssign', (('target', 'op', 'value',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('AnnAssign', (('target', 'annotation', 'value', 'simple',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Print', (('dest', 'values', 'nl',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('For', (('target', 'iter', 'body', 'orelse', 'type_comment'),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('AsyncFor', (('target', 'iter', 'body', 'orelse', 'type_comment'),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('While', (('test', 'body', 'orelse',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('If', (('test', 'body', 'orelse',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('With', (('items', 'body', 'type_comment'),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('AsyncWith', (('items', 'body', 'type_comment'),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Raise', (('exc', 'cause',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Try', (('body', 'handlers', 'orelse', 'finalbody',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Assert', (('test', 'msg',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Import', (('names',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('ImportFrom', (('module', 'names', 'level',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Exec', (('body', 'globals', 'locals',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Global', (('names',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Nonlocal', (('names',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Expr', (('value',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Pass', ((), ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Break', ((), ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
('Continue', ((),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(stmt,))),
|
|
|
|
# expr
|
|
|
|
('BoolOp', (('op', 'values',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('BinOp', (('left', 'op', 'right',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('UnaryOp', (('op', 'operand',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Lambda', (('args', 'body',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('IfExp', (('test', 'body', 'orelse',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Dict', (('keys', 'values',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Set', (('elts',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('ListComp', (('elt', 'generators',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('SetComp', (('elt', 'generators',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('DictComp', (('key', 'value', 'generators',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('GeneratorExp', (('elt', 'generators',),
|
|
('lineno', 'col_offset',
|
|
'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Await', (('value',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Yield', (('value',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('YieldFrom', (('value',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Compare', (('left', 'ops', 'comparators',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Call', (('func', 'args', 'keywords',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Repr', (('value',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('FormattedValue', (('value', 'conversion', 'format_spec',),
|
|
('lineno', 'col_offset',
|
|
'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('JoinedStr', (('values',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Constant', (('value', 'kind'),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Attribute', (('value', 'attr', 'ctx',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Subscript', (('value', 'slice', 'ctx',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Starred', (('value', 'ctx',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Name', (('id', 'ctx', 'annotation', 'type_comment'),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('List', (('elts', 'ctx',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
('Tuple', (('elts', 'ctx',),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(expr,))),
|
|
|
|
# expr_context
|
|
('Load', ((), (), (expr_context,))),
|
|
('Store', ((), (), (expr_context,))),
|
|
('Del', ((), (), (expr_context,))),
|
|
('AugLoad', ((), (), (expr_context,))),
|
|
('AugStore', ((), (), (expr_context,))),
|
|
('Param', ((), (), (expr_context,))),
|
|
|
|
# slice
|
|
('Slice', (('lower', 'upper', 'step'),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
|
|
(slice,))),
|
|
|
|
# boolop
|
|
('And', ((), (), (boolop,))),
|
|
('Or', ((), (), (boolop,))),
|
|
|
|
# operator
|
|
('Add', ((), (), (operator,))),
|
|
('Sub', ((), (), (operator,))),
|
|
('Mult', ((), (), (operator,))),
|
|
('MatMult', ((), (), (operator,))),
|
|
('Div', ((), (), (operator,))),
|
|
('Mod', ((), (), (operator,))),
|
|
('Pow', ((), (), (operator,))),
|
|
('LShift', ((), (), (operator,))),
|
|
('RShift', ((), (), (operator,))),
|
|
('BitOr', ((), (), (operator,))),
|
|
('BitXor', ((), (), (operator,))),
|
|
('BitAnd', ((), (), (operator,))),
|
|
('FloorDiv', ((), (), (operator,))),
|
|
|
|
# unaryop
|
|
('Invert', ((), (), (unaryop, AST,))),
|
|
('Not', ((), (), (unaryop, AST,))),
|
|
('UAdd', ((), (), (unaryop, AST,))),
|
|
('USub', ((), (), (unaryop, AST,))),
|
|
|
|
# cmpop
|
|
('Eq', ((), (), (cmpop,))),
|
|
('NotEq', ((), (), (cmpop,))),
|
|
('Lt', ((), (), (cmpop,))),
|
|
('LtE', ((), (), (cmpop,))),
|
|
('Gt', ((), (), (cmpop,))),
|
|
('GtE', ((), (), (cmpop,))),
|
|
('Is', ((), (), (cmpop,))),
|
|
('IsNot', ((), (), (cmpop,))),
|
|
('In', ((), (), (cmpop,))),
|
|
('NotIn', ((), (), (cmpop,))),
|
|
|
|
# comprehension
|
|
('comprehension', (('target', 'iter', 'ifs', 'is_async'), (), (AST,))),
|
|
|
|
# excepthandler
|
|
('ExceptHandler', (('type', 'name', 'body'),
|
|
('lineno', 'col_offset',
|
|
'end_lineno', 'end_col_offset'),
|
|
(excepthandler,))),
|
|
|
|
# arguments
|
|
('arguments', (('args', 'posonlyargs', 'vararg', 'kwonlyargs',
|
|
'kw_defaults', 'kwarg', 'defaults'), (), (AST,))),
|
|
|
|
# keyword
|
|
('keyword', (('arg', 'value'),
|
|
('lineno', 'col_offset', 'end_lineno', 'end_col_offset'),
|
|
(AST,))),
|
|
|
|
# alias
|
|
('alias', (('name', 'asname'), (), (AST,))),
|
|
|
|
# withitem
|
|
('withitem', (('context_expr', 'optional_vars'), (), (AST,))),
|
|
|
|
# type_ignore
|
|
('type_ignore', ((), ('lineno', 'tag'), (TypeIgnore,))),
|
|
)
|
|
|
|
for name, descr in _nodes:
|
|
_make_node(name, *descr)
|
|
|
|
if _sys.version_info.major == 2:
|
|
from .ast2 import ast_to_gast, gast_to_ast
|
|
if _sys.version_info.major == 3:
|
|
from .ast3 import ast_to_gast, gast_to_ast
|
|
|
|
|
|
def parse(*args, **kwargs):
|
|
return ast_to_gast(_ast.parse(*args, **kwargs))
|
|
|
|
|
|
def literal_eval(node_or_string):
|
|
if isinstance(node_or_string, AST):
|
|
node_or_string = gast_to_ast(node_or_string)
|
|
return _ast.literal_eval(node_or_string)
|
|
|
|
|
|
def get_docstring(node, clean=True):
|
|
if not isinstance(node, (FunctionDef, ClassDef, Module)):
|
|
raise TypeError("%r can't have docstrings" % node.__class__.__name__)
|
|
if node.body and isinstance(node.body[0], Expr) and \
|
|
isinstance(node.body[0].value, Constant):
|
|
if clean:
|
|
import inspect
|
|
holder = node.body[0].value
|
|
return inspect.cleandoc(getattr(holder, holder._fields[0]))
|
|
return node.body[0].value.s
|
|
|
|
|
|
# the following are directly imported from python3.8's Lib/ast.py #
|
|
|
|
def copy_location(new_node, old_node):
|
|
"""
|
|
Copy source location (`lineno`, `col_offset`, `end_lineno`, and
|
|
`end_col_offset` attributes) from *old_node* to *new_node* if possible,
|
|
and return *new_node*.
|
|
"""
|
|
for attr in 'lineno', 'col_offset', 'end_lineno', 'end_col_offset':
|
|
if attr in old_node._attributes and attr in new_node._attributes \
|
|
and hasattr(old_node, attr):
|
|
setattr(new_node, attr, getattr(old_node, attr))
|
|
return new_node
|
|
|
|
|
|
def fix_missing_locations(node):
|
|
"""
|
|
When you compile a node tree with compile(), the compiler expects lineno
|
|
and col_offset attributes for every node that supports them. This is
|
|
rather tedious to fill in for generated nodes, so this helper adds these
|
|
attributes recursively where not already set, by setting them to the values
|
|
of the parent node. It works recursively starting at *node*.
|
|
"""
|
|
def _fix(node, lineno, col_offset, end_lineno, end_col_offset):
|
|
if 'lineno' in node._attributes:
|
|
if not hasattr(node, 'lineno'):
|
|
node.lineno = lineno
|
|
else:
|
|
lineno = node.lineno
|
|
if 'end_lineno' in node._attributes:
|
|
if not hasattr(node, 'end_lineno'):
|
|
node.end_lineno = end_lineno
|
|
else:
|
|
end_lineno = node.end_lineno
|
|
if 'col_offset' in node._attributes:
|
|
if not hasattr(node, 'col_offset'):
|
|
node.col_offset = col_offset
|
|
else:
|
|
col_offset = node.col_offset
|
|
if 'end_col_offset' in node._attributes:
|
|
if not hasattr(node, 'end_col_offset'):
|
|
node.end_col_offset = end_col_offset
|
|
else:
|
|
end_col_offset = node.end_col_offset
|
|
for child in iter_child_nodes(node):
|
|
_fix(child, lineno, col_offset, end_lineno, end_col_offset)
|
|
_fix(node, 1, 0, 1, 0)
|
|
return node
|
|
|
|
|
|
def increment_lineno(node, n=1):
|
|
"""
|
|
Increment the line number and end line number of each node in the tree
|
|
starting at *node* by *n*. This is useful to "move code" to a different
|
|
location in a file.
|
|
"""
|
|
for child in walk(node):
|
|
if 'lineno' in child._attributes:
|
|
child.lineno = (getattr(child, 'lineno', 0) or 0) + n
|
|
if 'end_lineno' in child._attributes:
|
|
child.end_lineno = (getattr(child, 'end_lineno', 0) or 0) + n
|
|
return node
|