import ast,_ast
from collections import namedtuple
# List of classes for our internal AST
boolOp_e = namedtuple('boolOp_e',['op','args'])
binaryOp_e = namedtuple('binaryOp_e',['op','args'])
unaryOp_e = namedtuple('unaryOp_e',['op','arg'])
lambda_e = namedtuple('lambda_e',['args','body'])
if_e = namedtuple('if_e',['test','then','or_else'])
listComp_e = namedtuple('listComp_e',['expr','generators'])
setComp_e = namedtuple('setComp_e',['expr','generators'])
dictComp_e = namedtuple('dictComp_e',['key','value','generators'])
generatorExpr_e = namedtuple('generatorExpr_e',['expr','generators'])
compare_e = namedtuple('compare_e',['left','ops','comparators'])
call_e = namedtuple('call_e',['func','args','kwargs','starargs'])
num_literal = namedtuple('num_literal',['value'])
str_literal = namedtuple('str_literal',['value'])
bool_literal = namedtuple('bool_literal',['value'])
none_literal = namedtuple('none_literal',[])
name_e = namedtuple('name_e',['id'])
set_e = namedtuple('set_e',['values'])
list_e = namedtuple('list_e',['values'])
tuple_e = namedtuple('tuple_e',['values'])
dict_e = namedtuple('dict_e',['keys','values'])
attribute_e = namedtuple('attribute_e',['value','attribute'])
subscript_e = namedtuple('subscript_e',['value','lower','upper','step'])
subscript_ind_e = namedtuple('subscript_e',['value','index'])
slice_e = namedtuple('slice_e',['lower','upper','step'])
index_e = namedtuple('index_e',['value'])
comprehension_e = namedtuple('comprehension',['target','iter','ifs'])
# Mapping from Python's AST into the operators of our internal AST
opMap = {_ast.Add:'+', _ast.Sub:'-', _ast.UAdd:'+', _ast.USub:'-', _ast.Mult:'*', _ast.Div:'/',
_ast.Mod:'%', _ast.Pow:'**', _ast.LShift:'<<', ast.RShift:'>>',
_ast.BitOr:'???', _ast.BitXor:'???', _ast.BitAnd:'???', _ast.FloorDiv:'//',
_ast.Invert:'???', _ast.Not:'not', _ast.Eq:'==', _ast.NotEq:'!=',
_ast.Lt:'<', _ast.Gt:'>', _ast.LtE:'<=', _ast.GtE:'>=', _ast.Is:'is',
_ast.IsNot:'isnot', _ast.In:'in', _ast.NotIn:"notin", _ast.And:'and',
_ast.Or:'or'}
def isStarArg(x):
try:
return isinstance(x,_ast.Starred)
except:
return False
# Convert Python AST into our internal AST
def convert_ast(a):
if isinstance(a,_ast.BoolOp):
return boolOp_e(opMap[type(a.op)],
[convert_ast(x) for x in a.values])
elif isinstance(a,_ast.BinOp):
return binaryOp_e(opMap[type(a.op)],
[convert_ast(a.left), convert_ast(a.right)])
elif isinstance(a,_ast.UnaryOp):
return unaryOp_e(opMap[type(a.op)], convert_ast(a.operand))
elif isinstance(a,_ast.Lambda):
return lambda_e(convert_ast(a.args), convert_ast(a.body))
elif isinstance(a,_ast.IfExp):
return if_e(convert_ast(a.test),convert_ast(a.body), convert_ast(a.orelse))
elif isinstance(a,_ast.Attribute):
return attribute_e(convert_ast(a.value), name_e(a.attr))
elif isinstance(a,_ast.Subscript):
if isinstance(a.slice,ast.Index):
return subscript_ind_e(convert_ast(a.value), convert_ast(a.slice.value))
else:
return subscript_e(convert_ast(a.value),
convert_ast(a.slice.lower),
convert_ast(a.slice.upper),
convert_ast(a.slice.step))
elif isinstance(a,_ast.Compare):
return compare_e(convert_ast(a.left),
[opMap[type(x)] for x in a.ops],
[convert_ast(x) for x in a.comparators])
elif isinstance(a,_ast.Call):
return call_e(convert_ast(a.func),
[convert_ast(x) for x in a.args if not isStarArg(x)],
[{k.arg : convert_ast(k.value)} for k in a.keywords],
[convert_ast(x.value) for x in a.args if isStarArg(x)]
)
elif isinstance(a,_ast.ListComp):
return listComp_e(convert_ast(a.elt), [convert_ast(x) for x in a.generators])
elif isinstance(a,_ast.SetComp):
return setComp_e(convert_ast(a.elt), [convert_ast(x) for x in a.generators])
elif isinstance(a,_ast.DictComp):
return dictComp_e(convert_ast(a.key), convert_ast(a.value), [convert_ast(x) for x in a.generators])
elif isinstance(a,_ast.comprehension):
return comprehension_e(convert_ast(a.target), convert_ast(a.iter), convert_ast(a.ifs))
elif isinstance(a,_ast.List):
return list_e([convert_ast(x) for x in a.elts])
elif isinstance(a,_ast.Tuple):
return tuple_e([convert_ast(x) for x in a.elts])
elif isinstance(a,_ast.Set):
return set_e([convert_ast(x) for x in a.elts])
elif isinstance(a,_ast.Dict):
return dict_e([convert_ast(x) for x in a.keys],[convert_ast(x) for x in a.values])
elif isinstance(a,_ast.Str):
return str_literal(a.s)
elif isinstance(a,_ast.Num):
return num_literal(a.n)
elif isinstance(a,_ast.Name):
return name_e(a.id)
elif isinstance(a,_ast.NameConstant):
if type(a.value) == bool:
return bool_literal(a.value)
else:
return none_literal()
all_ast_types = [boolOp_e,
binaryOp_e,
unaryOp_e,
lambda_e,
if_e,
listComp_e,
setComp_e,
dictComp_e,
generatorExpr_e,
compare_e,
call_e,
num_literal,
str_literal,
bool_literal,
none_literal,
name_e,
set_e,
list_e,
tuple_e,
dict_e,
attribute_e,
subscript_ind_e,
subscript_e,
slice_e,
index_e,
comprehension_e]
literal_ast_types = [num_literal,
str_literal,
bool_literal,
none_literal]
comprehension_types = [listComp_e,
setComp_e,
dictComp_e,
generatorExpr_e]
def is_ast(t):
return type(t) in all_ast_types
def is_literal(t):
return type(t) in literal_ast_types
def is_comprehension(t):
return type(t) in comprehension_types
# Get all variables used in the AST expression
def get_all_vars(a):
if is_ast(a) and not is_literal(a):
if type(a) == name_e:
return {a.id}
if is_comprehension(a):
vs = list(get_all_vars(a.expr)) + [v for g in a.generators for v in get_all_vars(g)]
vs = set(vs) - {v for g in a.generators for v in get_all_vars(g.target)}
return vs
if type(a)==call_e:
# We need a special case for the make_pql_tuple
if isinstance(a.func, name_e) and a.func.id == 'make_pql_tuple':
t = a.args[0].values
vs = [v for x in t for v in get_all_vars(get_ast(x.values[0].value))]
return set(vs)
# And we need a special case for nested queries also. However, instead of digging
# into the nested query, we just return an impossible variable, so that nothing can
# satisfy its dependency.
if isinstance(a.func, name_e) and a.func.id == 'PyQuery':
return {"#nested_query"}
else:
vs = [v for x in a[1:] for y in x for v in get_all_vars(y) ]
return set(vs)
if type(a)==attribute_e:
return get_all_vars(a.value)
retvars = set()
for x in a:
if is_ast(x):
retvars = retvars.union(get_all_vars(x))
elif type(x)==list:
for y in x:
if is_ast(y):
retvars = retvars.union(get_all_vars(y))
return retvars
return set()
# Replace variables inside an expression accorind to the table
# of mappings
def replace_vars(a,table):
if is_ast(a) and not is_literal(a):
if type(a) == name_e:
if a.id in table:
return table[a.id]
else:
return a
if type(a)==call_e:
return call_e(a.func, *([replace_vars(x,table) for x in a[1:]]))
if type(a)==attribute_e:
return attribute_e(replace_vars(a.value,table),a.attribute)
for (i,x) in enumerate(a):
if is_ast(x):
a = a._replace(**{a._fields[i]:replace_vars(x,table)})
elif type(x)==list:
for (j,y) in enumerate(x):
x[j] = replace_vars(y,table)
return a
# Visit all nodes in the AST
def visit(a):
if is_ast(a):
yield a
for x in a:
if is_ast(x):
for y in visit(x):
yield y
elif type(x)==list:
for y in x:
for z in visit(y):
yield z
# Compile into Python's AST and convert into our internal AST format
def get_ast(expr):
return convert_ast(compile(expr, '