def visit_stmt(self, node):
# type: (ast.stmt) -> ast.With
"""
Every statement in the original code becomes:
with _treetrace_hidden_with_stmt(_tree_index):
<statement>
where the _treetrace_hidden_with_stmt function is the the corresponding method with the
TreeTracerBase and traced_file arguments already filled in (see _trace_methods_dict)
"""
context_expr = self._create_simple_marker_call(
super(_NodeVisitor, self).generic_visit(node),
TreeTracerBase._treetrace_hidden_with_stmt)
if PY3:
wrapped = ast.With(
items=[ast.withitem(context_expr=context_expr)],
body=[node],
)
else:
wrapped = ast.With(
context_expr=context_expr,
body=[node],
)
ast.copy_location(wrapped, node)
ast.fix_missing_locations(wrapped)
return wrapped
def visitWith(self, n, *args):
body = self.dispatch(n.body, *args)
if flags.PY_VERSION == 3 and flags.PY3_VERSION >= 3:
items = [self.dispatch(item, *args) for item in n.items]
# prots = [itm.retic_protector for itm in items if itm.retic_protector]
return ast.With(items=items, body=body) #prots + body)
else:
cexpr = self.dispatch(n.context_expr, *args)
optvars = self.dispatch(n.optional_vars, *args)
if optvars:
if isinstance(optvars, ast.Name):
cexpr = retic_ast.Check(value=cexpr,
type=optvars.retic_type,
lineno=cexpr.lineno,
col_offset=cexpr.col_offset)
elif isinstance(target, ast.Starred):
raise exc.UnimplementedException(
'Assignment checks against Starred')
elif isinstance(target, ast.List):
raise exc.UnimplementedException(
'Assignment checks against List')
elif isinstance(target, ast.Tuple):
raise exc.UnimplementedException(
'Assignment checks against Tuple')
# prot = self.handlewithitem(optvars)
# if prot:
# body = [prot] + body
return ast.With(context_expr=cexpr,
optional_vars=optvars,
body=body)
def test_deeply_nested(self):
with_C = ast.With([self.C_clause], [ast.Pass()])
with_B = ast.With([self.B_clause], [with_C])
with_A = ast.With([self.A_clause], [with_B])
new_ast = self.transform.visit(with_A)
expect = ast.With([self.A_clause, self.B_clause, self.C_clause],
[ast.Pass()])
self.assertEqual(ast.dump(new_ast), ast.dump(expect))
def test_with(self):
p = ast.Pass()
self.stmt(ast.With([], [p]), "empty items on With")
i = ast.withitem(ast.Num(3), None)
self.stmt(ast.With([i], []), "empty body on With")
i = ast.withitem(ast.Name("x", ast.Store()), None)
self.stmt(ast.With([i], [p]), "must have Load context")
i = ast.withitem(ast.Num(3), ast.Name("x", ast.Load()))
self.stmt(ast.With([i], [p]), "must have Store context")
def _make_with_node(self, with_expr, body):
if self.HAS_WITHITEM:
return ast.With(
items=[ast.withitem(
context_expr=with_expr,
optional_vars=None)],
body=body)
else:
return ast.With(
context_expr=with_expr,
optional_vars=None,
body=body)
def visitWith(self, n, *args):
body = self.dispatch_statements(n.body, *args)
if flags.PY_VERSION == 3 and flags.PY3_VERSION >= 3:
items = [self.dispatch(item, *args) for item in n.items]
return ast.With(items=items, body=body)
else:
context = self.dispatch(n.context_expr, *args)
optional_vars = self.dispatch(n.optional_vars, *
args) if n.optional_vars else None
return ast.With(context_expr=context,
optional_vars=optional_vars,
body=body)
def generate_with(max_depth=None):
num_items = random.randrange(1, 3)
items = [
_generate_with_item(max_depth=max_depth) for _ in range(num_items)
]
body = generate_block(max_depth=max_depth)
return ast.With(items, body)
def make_async(cursor_trail, tree):
definition = core_logic.get_node_at_cursor(cursor_trail, tree)
if type(definition) == ast.FunctionDef:
return ast.AsyncFunctionDef(name=definition.name,
args=definition.args,
body=definition.body,
decorator_list=definition.decorator_list,
returns=definition.returns,
type_comment=definition.type_comment)
if type(definition) == ast.AsyncFunctionDef:
return ast.FunctionDef(name=definition.name,
args=definition.args,
body=definition.body,
decorator_list=definition.decorator_list,
returns=definition.returns,
type_comment=definition.type_comment)
if type(definition) == ast.With:
return ast.AsyncWith(items=definition.items,
body=definition.body,
type_comment=definition.type_comment)
if type(definition) == ast.AsyncWith:
return ast.With(items=definition.items,
body=definition.body,
type_comment=definition.type_comment)
if isinstance(definition, ast.comprehension):
# Toggle the async
definition.is_async = 0 if definition.is_async == 1 else 1
return definition
def visit_With(self, node):
new_node = ast.With(
self._visit(node.items),
self._visit(node.body),
)
ast.copy_location(new_node, node)
return new_node
def compile_with_expression(self, expr):
expr.pop(0) # with
args = expr.pop(0)
if len(args) > 2 or len(args) < 1:
raise HyTypeError(expr, "with needs [arg (expr)] or [(expr)]")
args.reverse()
ctx = self.compile(args.pop(0))
thing = None
if args != []:
thing = self._storeize(self.compile(args.pop(0)))
ret = ast.With(context_expr=ctx,
lineno=expr.start_line,
col_offset=expr.start_column,
optional_vars=thing,
body=self._code_branch([self.compile(x) for x in expr],
expr.start_line,
expr.start_column))
if sys.version_info[0] >= 3 and sys.version_info[1] >= 3:
ret.items = [ast.withitem(context_expr=ctx, optional_vars=thing)]
return ret
def visit_With(self, node):
new_body = self.if_exists(node.body, self.visit_list)
new_with_items = self.if_exists(node.items, self.visit_list)
return_list = self.new_stmts + [
ast.copy_location(ast.With(new_with_items, new_body), node)
]
self.new_stmts.clear()
return return_list
def visit_With(self, node):
new_node = ast.With(
self._visit(node.items[0].context_expr),
self._visit(node.items[0].optional_vars),
self._visit(node.body)
)
ast.copy_location(new_node, node)
return new_node
def wrap_in_with(self, name, val, body):
name_expr = ast.Str(name)
val_expr = ast.Num(val)
args = [name_expr, val_expr]
scope_expr = ast.Call(func=ast.Name(id='scope', ctx=ast.Load()),
args=args,
keywords=[])
return [ast.With(items=[ast.withitem(scope_expr, None)], body=body)]
def explode_with(stmt: ast.With) -> t.List[ast.stmt]:
if stmt.items:
res: t.List[ast.stmt] = []
first: ast.withitem = stmt.items.pop(0)
first.context_expr = explode_expr(first.context_expr, res)
res.append(ast.With(items=[first], body=explode_with(stmt)))
return res
else:
return explode_stmts(stmt.body)
def visit_stmt(self, node):
context_expr = _create_simple_marker_call(
super(_NodeVisitor, self).generic_visit(node),
TreeTracerBase._treetrace_hidden_with_stmt)
if PY3:
wrapped = ast.With(
items=[ast.withitem(context_expr=context_expr)],
body=[node],
)
else:
wrapped = ast.With(
context_expr=context_expr,
body=[node],
)
ast.copy_location(wrapped, node)
ast.fix_missing_locations(wrapped)
return wrapped
def simplify_with(stmt: ast.With) -> t.List[ast.stmt]:
if stmt.items:
res: t.List[ast.stmt] = []
first: ast.withitem = stmt.items.pop(0)
first.context_expr = SimplifyExpr(res).visit(first.context_expr)
res.append(ast.With(items=[first], body=simplify_with(stmt)))
return res
else:
return simplify_stmts(stmt.body)
def visitWith(self, n, *args):
body = self.dispatch(n.body, *args)
if flags.PY_VERSION == 3 and flags.PY3_VERSION >= 3:
items = [self.dispatch(item, *args) for item in n.items]
prots = [
itm.retic_protector for itm in items if itm.retic_protector
]
return ast.With(items=items, body=prots + body)
else:
cexpr = self.dispatch(n.context_expr, *args)
optvars = self.dispatch(n.optional_vars, *args)
prot = self.handlewithitem(optvars)
if prot:
body = [prot] + body
return ast.With(context_expr=cexpr,
optional_vars=optvars,
body=body)
def visit_With(self, node: With, *args, **kwargs) -> C.With:
items = self.visit(node.items, *args, **kwargs)
body = self.visit(node.body, *args, **kwargs)
type_comment = self.visit(node.type_comment, *args, **kwargs)
return C.With(
items=items,
body=body,
type_comment=type_comment,
)
def wrap_in_method(self, body):
method_name_expr = ast.Str(methods[-1])
args = [method_name_expr]
scope_expr = ast.Call(func=ast.Name(id='method__', ctx=ast.Load()),
args=args,
keywords=[])
return [
ast.With(
items=[ast.withitem(scope_expr, ast.Name(id='_method__'))],
body=body)
]
def wrap_in_outer(self, name, counter, node):
#return node # <- to disable
name_expr = ast.Str(name)
counter_expr = ast.Num(counter)
args = [name_expr, counter_expr, ast.Name(id=TaintName)]
scope_expr = ast.Call(func=ast.Name(id=TaintedScope, ctx=ast.Load()),
args=args,
keywords=[])
return ast.With(
items=[ast.withitem(scope_expr, ast.Name(id=TaintName))],
body=[node])
def test_With(self):
# with A: pass
A = ast.Name('A', ast.Load())
A_clause = ast.withitem(A, None)
with_A = ast.With([A_clause], [ast.Pass()])
self.verify(with_A, 'with A:pass')
# with A as a: pass
a = ast.Name('a', ast.Store())
a_clause = ast.withitem(A, a)
with_a = ast.With([a_clause], [ast.Pass()])
self.verify(with_a, 'with A as a:pass')
# with A as A, B: pass
B = ast.Name('B', ast.Load())
B_clause = ast.withitem(B, None)
with_B = ast.With([a_clause, B_clause], [ast.Pass()])
self.verify(with_B, 'with A as a,B:pass')
# with A as A, B as b: pass
b = ast.Name('b', ast.Store())
b_clause = ast.withitem(B, b)
with_b = ast.With([a_clause, b_clause], [ast.Pass()])
self.verify(with_b, 'with A as a,B as b:pass')
def wrap_in_method(self, body, args):
method_name_expr = ast.Str(methods[-1])
my_args = ast.List(args.args, ast.Load())
args = [method_name_expr, my_args]
scope_expr = ast.Call(func=ast.Name(id=TaintedMethod, ctx=ast.Load()),
args=args,
keywords=[])
# we expect the method__ to push in the taint in the beginning, and pop out when it is done.
return [
ast.With(items=[ast.withitem(scope_expr, ast.Name(id=TaintName))],
body=body)
]
def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST:
with_item = ast.withitem(context_expr=ast.Call(func=ast.Name(
id='SoftAssertions', ctx=ast.Load()),
args=[],
keywords=[]),
optional_vars=ast.Name(id='ctx',
ctx=ast.Store()))
with_stmt = ast.With(items=[with_item], body=node.body)
node.body = [with_stmt]
new_node = fix_node(node)
return self.generic_visit(new_node)
def make_with():
"""With(withitem* items, stmt* body, string? type_comment)"""
"""withitem = (expr context_expr, expr? optional_vars)"""
return ast.With(
items=[
ast.withitem(context_expr=ast.Name(id="context_manager",
ctx=ast.Load()),
optional_vars=ast.Name(id="f", ctx=ast.Store()))
],
body=[make_statement()],
type_comment=None,
)
def wrap_in_outer(self, name, can_empty, counter, node):
name_expr = ast.Str(name)
can_empty_expr = ast.Str(can_empty)
counter_expr = ast.Num(counter)
method_id = ast.Name(id='_method__')
args = [name_expr, counter_expr, method_id, can_empty_expr]
scope_expr = ast.Call(func=ast.Name(id='stack__', ctx=ast.Load()),
args=args,
keywords=[])
return ast.With(items=[
ast.withitem(scope_expr,
ast.Name(id='%s_%d_stack__' % (name, counter)))
],
body=[node])
def section_stmt_action(s, loc, tokens):
body = tokens[1][:]
if PY3:
return ast.With(items=[
ast.withitem(context_expr=ast.Call(func=ast.Name(id='_section',
ctx=ast.Load()),
args=[tokens[0]],
keywords=[],
starargs=None,
kwargs=None),
optional_vars=None)
],
body=body)
else:
return ast.With(context_expr=ast.Call(func=ast.Name(id='_section',
ctx=ast.Load()),
args=[tokens[0]],
keywords=[],
starargs=None,
kwargs=None),
optional_vars=None,
body=body)
def visit_While(self, tree_node):
global while_counter
while_counter += 1
test = tree_node.test
body = tree_node.body
assert not tree_node.orelse
self.generic_visit(tree_node)
name = ast.Str('while_%d %s' % (while_counter, method_name))
val = ast.Num(0)
scope_expr = ast.Call(func=ast.Name(id='scope', ctx=ast.Load()),
args=[name, val],
keywords=[])
tree_node.body = [
ast.With(items=[ast.withitem(scope_expr, None)], body=body)
]
return tree_node
def wrap_in_inner(self, name, counter, val, body):
val_expr = ast.Num(val)
stack_iter = ast.Name(id='%s_%d_stack__' % (name, counter))
method_id = ast.Name(id='_method__')
args = [val_expr, stack_iter, method_id]
scope_expr = ast.Call(func=ast.Name(id='scope__', ctx=ast.Load()),
args=args,
keywords=[])
return [
ast.With(items=[
ast.withitem(
scope_expr,
ast.Name(id='%s_%d_%d_scope__' % (name, counter, val)))
],
body=body)
]
def _rewrite_if(tree, var_name=None, **kw_args):
# TODO refactor into a _rewrite_switch and a _rewrite_if
"""
Rewrite if statements to treat pattern matches as boolean expressions.
Recall that normally a pattern match is a statement which will throw a
PatternMatchException if the match fails. We can therefore use try-blocks
to produce the desired branching behavior.
var_name is an optional parameter used for rewriting switch statements. If
present, it will transform predicates which are expressions into pattern
matches.
"""
# with q as rewritten:
# try:
# with matching:
# u%(matchPattern)
# u%(successBody)
# except PatternMatchException:
# u%(_maybe_rewrite_if(failBody))
# return rewritten
if not isinstance(tree, ast.If):
return tree
if var_name:
tree.test = ast.BinOp(tree.test, ast.LShift(),
ast.Name(var_name, ast.Load()))
elif not (isinstance(tree.test, ast.BinOp) and \
isinstance(tree.test.op, ast.LShift)):
return tree
handler = ast.ExceptHandler(hq[PatternMatchException], None, tree.orelse)
try_stmt = ast.Try(tree.body, [handler], [], [])
macroed_match = ast.With(
[ast.withitem(ast.Name('_matching', ast.Load()), None)],
[ast.Expr(tree.test)])
try_stmt.body = [macroed_match] + try_stmt.body
if len(handler.body) == 1: # (== tree.orelse)
# Might be an elif
handler.body = [_rewrite_if(handler.body[0], var_name)]
elif not handler.body:
handler.body = [ast.Pass()]
return try_stmt
def util_json_builder(self):
"""
Return: ast sub-tree of a function to read and parse json file:
with io.open(path, mode='r', encoding='utf-8') as f:
return json.loads(f.read())
"""
io_open = ast_mod.Call(
func=ast_mod.Attribute(value=ast_name('io'),
attr=ast_name('open')),
args=[ast_name('path')],
keywords=[
ast_mod.keyword(arg='mode', value=ast_mod.Constant(value='r')),
ast_mod.keyword(arg='encoding',
value=ast_mod.Constant(value='utf-8'))
])
f_read = ast_mod.Call(ast_mod.Attribute(value=ast_name('f'),
attr=ast_name('read')),
args=[],
keywords=[])
return_node = ast_mod.Return(value=ast_mod.Call(func=ast_mod.Attribute(
value=ast_name('json'), attr=ast_name('loads')),
args=[f_read],
keywords=[]))
body = [
ast_mod.With(items=[
ast_mod.withitem(context_expr=io_open,
optional_vars=ast_name('f'))
],
body=[ast_mod.Expr(value=return_node)])
]
func = ast_mod.FunctionDef(name="util_load_json",
args=ast_mod.arguments(
posonlyargs=[],
args=[ast_name('path')],
vararg=None,
kwonlyargs=[],
kw_defaults=[],
kwarg=None,
defaults=[]),
body=body,
decorator_list=[],
returns=None)
return func
