def visit_Expr(self_t, node):
if u.is_set_to(node) and isinstance(node.value.args[0],
ast.Num):
rhs = node.value.args[0].n
var_name = node.value.func.value.id
var_size = self_t.size_visitor.get_var_size(var_name)
expr = u.stmt_from_str("tpt.one_hot(%s, %s)" %
(rhs, var_size))
node.value.args[0] = expr.value
return node
def vars_defined_in_all_cases(self, node):
if isinstance(node, ast.If):
vars_defined = None
for if_node in u.ifs_in_elif_block(node):
cur_vars_defined = self.vars_defined_in_all_cases(if_node.body)
if vars_defined is None:
vars_defined = cur_vars_defined
else:
vars_defined = vars_defined & cur_vars_defined
elif isinstance(node, list):
vars_defined = set()
for stmt_node in node:
vars_defined = vars_defined | self.vars_defined_in_all_cases(
stmt_node)
elif isinstance(node, ast.AST):
vars_defined = set()
if u.is_set_to(node):
vars_defined.add(node.value.func.value.id)
return vars_defined
def visit_Expr(self_t, node):
if u.is_set_to(node):
lhs = node.value.func.value
rhs = node.value.args[0]
node = ast.Assign(targets=[lhs], value=rhs)
return node
def visit_Expr(self_t, node):
if u.is_set_to(node):
name = node.value.func.value.id
self_t.set_vars.add(name)
return node