Spaces:
Running
Running
File size: 1,322 Bytes
6a86ad5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
from sympy.core.basic import Basic
from sympy.core.numbers import Integer
from sympy.core.singleton import S
from sympy.strategies.branch.traverse import top_down, sall
from sympy.strategies.branch.core import do_one, identity
def inc(x):
if isinstance(x, Integer):
yield x + 1
def test_top_down_easy():
expr = Basic(S(1), S(2))
expected = Basic(S(2), S(3))
brl = top_down(inc)
assert set(brl(expr)) == {expected}
def test_top_down_big_tree():
expr = Basic(S(1), Basic(S(2)), Basic(S(3), Basic(S(4)), S(5)))
expected = Basic(S(2), Basic(S(3)), Basic(S(4), Basic(S(5)), S(6)))
brl = top_down(inc)
assert set(brl(expr)) == {expected}
def test_top_down_harder_function():
def split5(x):
if x == 5:
yield x - 1
yield x + 1
expr = Basic(Basic(S(5), S(6)), S(1))
expected = {Basic(Basic(S(4), S(6)), S(1)), Basic(Basic(S(6), S(6)), S(1))}
brl = top_down(split5)
assert set(brl(expr)) == expected
def test_sall():
expr = Basic(S(1), S(2))
expected = Basic(S(2), S(3))
brl = sall(inc)
assert list(brl(expr)) == [expected]
expr = Basic(S(1), S(2), Basic(S(3), S(4)))
expected = Basic(S(2), S(3), Basic(S(3), S(4)))
brl = sall(do_one(inc, identity))
assert list(brl(expr)) == [expected]
|