def test_subtree():
exp1 = Add(Mul(Integer(1), Integer(2)), Div(Integer(3), Integer(4)))
tree1 = Tree(exp1)
tree2 = tree1.subtree(tree1.root_nid)
tree3 = tree1.subtree(tree1.root_nid, deep=True)
assert tree1.get_node(tree1.root_nid) == tree2.get_node(
tree2.root_nid) == tree3.get_node(tree3.root_nid)
assert set(tree1._nodes) == set(tree2._nodes)
assert not set(tree2._nodes).intersection(set(tree3._nodes))
def test_hash():
class A(Basic):
pass
assert hash(Basic(1, 2)) == hash(Basic(1, 2))
assert hash(Basic(1, Basic(2, 3))) == hash(Basic(1, Basic(2, 3)))
assert hash(A(1, 2)) != hash(Basic(1, 2))
expr1 = Add(
Mul(Integer(2), X),
Neg(Mul(Integer(3), Pow(X, Two))),
Neg(Integer(7)),
Mul(Integer(7), X),
Neg(Pow(X, Two)),
)
expr2 = Add(
Mul(Integer(2), X),
Neg(Mul(Integer(3), Pow(X, Two))),
Neg(Integer(7)),
Mul(Integer(7), X),
Neg(Pow(X, Two)),
)
assert hash(expr1) == hash(expr2)
def test_tree_string():
exp1 = Add(Mul(Integer(1), Integer(2)), Div(Integer(3), Symbol("x")))
tree = Tree(exp1)
tree_string_exp = (""
"Add\n"
"|-- Mul\n"
"| |-- Integer(1)\n"
"| +-- Integer(2)\n"
"+-- Div\n"
" |-- Integer(3)\n"
" +-- Symbol(x)\n")
assert tree.tree_string() == tree_string_exp
pass
assert A(1, 2) == A(1, 2)
assert B(1, 2) == B(1, 2)
assert A(1, 2) != B(1, 2)
def test_parent():
nested = Basic(Basic(Basic(1), 2), 3)
assert nested.parent is None
assert nested.args[0].parent is nested
assert nested.args[0].args[0].parent is nested.args[0]
@mark.parametrize(
"expr, types, expected_atoms",
[
(One, (), {One}),
(Add(One, Two, One), (), {One, Two}),
(Add(One, Two, One), (Symbol,), set()),
(Add(One, X), (Integer,), {One}),
(Add(One, X, Y), (Symbol,), {X, Y}),
(Mul(Add(One, X, Y), Two), (Add,), {Add(One, X, Y)}),
(Mul(Add(One, X, Y), Two), (Add, Symbol), {X, Y, Add(One, X, Y)}),
(Eq(X, Add(One, Two)), (Symbol,), {X}),
],
)
def test_atoms(expr, types, expected_atoms):
assert expr.atoms(*types) == expected_atoms
@fixture
def latex_printer():
return LatexPrinter()
@mark.parametrize(
"expression, latex",
[
(Add(One, Two), r"1+2"),
(Add(One, Two, Three), r"1+2+3"),
(Add(Two, Neg(One)), r"2-1"),
(Add(Three, Neg(Two), Neg(One)), r"3-2-1"),
(Add(Three, Neg(Add(Two, Neg(One)))), r"3-\left(2-1\right)"),
(Mul(Two, Three), r"2 \cdot 3"),
(Mul(Two, Three, Four), r"2 \cdot 3 \cdot 4"),
(Neg(One), r"-1"),
(Neg(Two), r"-2"),
(Add(Neg(Two), Neg(Four)), r"-2-4"),
(Add(Two, Neg(Neg(Four))), r"2-\left(-4\right)"),
(Mul(Two, Add(Two, Three)), r"2 \cdot \left(2+3\right)"),
(Div(Float(1.2), Float(0.1), print_format=PrintFormat.DIV_RATIONAL), r"\frac{1.2}{0.1}"),
(Div(Float(1.2), Float(0.1), print_format=PrintFormat.DIV_COLON), r"1.2 : 0.1"),
(Div(Float(1.2), Add(One, Two), print_format=PrintFormat.DIV_COLON), r"1.2 : \left(1+2\right)"),
(Div(Float(1.2), Mul(One, Two), print_format=PrintFormat.DIV_COLON), r"1.2 : \left(1 \cdot 2\right)"),
(Neg(Div(One, Two, print_format=PrintFormat.DIV_RATIONAL)), r"-\frac{1}{2}"),
(Mul(Add(Two, Three), Four), r"\left(2+3\right) \cdot 4"),
(Add(Two, Neg(Mul(Two, Three)), Four), r"2-2 \cdot 3+4"),
(Div(Two, One, print_format=PrintFormat.DIV_COLON), r"2 : 1"),
(
def test_init():
exp1 = Add(Mul(Integer(1), Integer(2)), Div(Integer(3), Integer(4)))
tree = Tree(exp1)
assert tree.size == 7
def test_leaves():
exp1 = Add(Mul(Integer(1), Integer(2)), Div(Integer(3), Symbol("x")))
tree = Tree(exp1)
assert (set(tree.leaves()) == set(tree.leaves(nid=tree.root_nid)) ==
{Integer(1), Integer(2),
Integer(3), Symbol("x")})