def testReplaceNodeWithExpr(self):
def f(x):
return 2 * x
def g(x):
return 3 * x
expr = tracers.make_expr(f, 5)
new_expr = tracers.make_expr(g, 10)
tracers.replace_node_with_expr(expr.expr_node, new_expr)
self.assertEqual(3 * 7, tracers.eval_expr(expr, {'x': 7}))
def testExtractSuperexprWithReplaceNode(self):
# NOTE(mattjj): this test shows an alternative way to implement, in effect,
# tracers.extract_superexpr just using tracers.replace_node_with_expr. The
# reason to have both is that one does in-place modification.
def f(x, y):
return 2 * x ** 2 + y
expr = tracers.make_expr(f, 1, 2)
node = expr.expr_node.parents[0].parents[0] # x ** 2
lookup_expr = tracers.make_expr(lambda x: x, 3, names=('x2',))
tracers.replace_node_with_expr(node, lookup_expr) # modify expr in-place
self.assertEqual(2 * 5 + 6, tracers.eval_expr(expr, {'x2': 5, 'y': 6}))
def testInlineExprAndReplace(self):
def f(x, y):
return 2 * x ** 2 + y
def g(z):
return 3 * z ** 3
expr = tracers.make_expr(f, 1, 2)
subexpr = tracers.make_expr(g, 3)
input_node = expr.expr_node.parents[0].parents[0] # x ** 2
output_node = expr.expr_node.parents[0] # 2 * x ** 2
new_expr = tracers.inline_expr(subexpr, {'z': input_node})
tracers.replace_node_with_expr(output_node, new_expr) # modify expr inplace
self.assertEqual(3 * 6 ** 6 + 7,
tracers.eval_expr(expr, {'x': 6, 'y': 7}))