<html><head><meta name="color-scheme" content="light dark"></head><body><pre style="word-wrap: break-word; white-space: pre-wrap;">from sympy import symbols, sin, Matrix, Interval, Piecewise, Sum, lambdify, \
                  Expr, sqrt
from sympy.testing.pytest import raises

from sympy.printing.tensorflow import TensorflowPrinter
from sympy.printing.lambdarepr import lambdarepr, LambdaPrinter, NumExprPrinter


x, y, z = symbols("x,y,z")
i, a, b = symbols("i,a,b")
j, c, d = symbols("j,c,d")


def test_basic():
    assert lambdarepr(x*y) == "x*y"
    assert lambdarepr(x + y) in ["y + x", "x + y"]
    assert lambdarepr(x**y) == "x**y"


def test_matrix():
    A = Matrix([[x, y], [y*x, z**2]])
    # assert lambdarepr(A) == "ImmutableDenseMatrix([[x, y], [x*y, z**2]])"
    # Test printing a Matrix that has an element that is printed differently
    # with the LambdaPrinter than in the StrPrinter.
    p = Piecewise((x, True), evaluate=False)
    A = Matrix([p])
    assert lambdarepr(A) == "ImmutableDenseMatrix([[((x))]])"


def test_piecewise():
    # In each case, test eval() the lambdarepr() to make sure there are a
    # correct number of parentheses. It will give a SyntaxError if there aren't.

    h = "lambda x: "

    p = Piecewise((x, True), evaluate=False)
    l = lambdarepr(p)
    eval(h + l)
    assert l == "((x))"

    p = Piecewise((x, x &lt; 0))
    l = lambdarepr(p)
    eval(h + l)
    assert l == "((x) if (x &lt; 0) else None)"

    p = Piecewise(
        (1, x &lt; 1),
        (2, x &lt; 2),
        (0, True)
    )
    l = lambdarepr(p)
    eval(h + l)
    assert l == "((1) if (x &lt; 1) else (2) if (x &lt; 2) else (0))"

    p = Piecewise(
        (1, x &lt; 1),
        (2, x &lt; 2),
    )
    l = lambdarepr(p)
    eval(h + l)
    assert l == "((1) if (x &lt; 1) else (2) if (x &lt; 2) else None)"

    p = Piecewise(
        (x, x &lt; 1),
        (x**2, Interval(3, 4, True, False).contains(x)),
        (0, True),
    )
    l = lambdarepr(p)
    eval(h + l)
    assert l == "((x) if (x &lt; 1) else (x**2) if (((x &lt;= 4)) and ((x &gt; 3))) else (0))"

    p = Piecewise(
        (x**2, x &lt; 0),
        (x, x &lt; 1),
        (2 - x, x &gt;= 1),
        (0, True), evaluate=False
    )
    l = lambdarepr(p)
    eval(h + l)
    assert l == "((x**2) if (x &lt; 0) else (x) if (x &lt; 1)"\
                                " else (2 - x) if (x &gt;= 1) else (0))"

    p = Piecewise(
        (x**2, x &lt; 0),
        (x, x &lt; 1),
        (2 - x, x &gt;= 1), evaluate=False
    )
    l = lambdarepr(p)
    eval(h + l)
    assert l == "((x**2) if (x &lt; 0) else (x) if (x &lt; 1)"\
                    " else (2 - x) if (x &gt;= 1) else None)"

    p = Piecewise(
        (1, x &gt;= 1),
        (2, x &gt;= 2),
        (3, x &gt;= 3),
        (4, x &gt;= 4),
        (5, x &gt;= 5),
        (6, True)
    )
    l = lambdarepr(p)
    eval(h + l)
    assert l == "((1) if (x &gt;= 1) else (2) if (x &gt;= 2) else (3) if (x &gt;= 3)"\
                        " else (4) if (x &gt;= 4) else (5) if (x &gt;= 5) else (6))"

    p = Piecewise(
        (1, x &lt;= 1),
        (2, x &lt;= 2),
        (3, x &lt;= 3),
        (4, x &lt;= 4),
        (5, x &lt;= 5),
        (6, True)
    )
    l = lambdarepr(p)
    eval(h + l)
    assert l == "((1) if (x &lt;= 1) else (2) if (x &lt;= 2) else (3) if (x &lt;= 3)"\
                            " else (4) if (x &lt;= 4) else (5) if (x &lt;= 5) else (6))"

    p = Piecewise(
        (1, x &gt; 1),
        (2, x &gt; 2),
        (3, x &gt; 3),
        (4, x &gt; 4),
        (5, x &gt; 5),
        (6, True)
    )
    l = lambdarepr(p)
    eval(h + l)
    assert l =="((1) if (x &gt; 1) else (2) if (x &gt; 2) else (3) if (x &gt; 3)"\
                            " else (4) if (x &gt; 4) else (5) if (x &gt; 5) else (6))"

    p = Piecewise(
        (1, x &lt; 1),
        (2, x &lt; 2),
        (3, x &lt; 3),
        (4, x &lt; 4),
        (5, x &lt; 5),
        (6, True)
    )
    l = lambdarepr(p)
    eval(h + l)
    assert l == "((1) if (x &lt; 1) else (2) if (x &lt; 2) else (3) if (x &lt; 3)"\
                            " else (4) if (x &lt; 4) else (5) if (x &lt; 5) else (6))"

    p = Piecewise(
        (Piecewise(
            (1, x &gt; 0),
            (2, True)
        ), y &gt; 0),
        (3, True)
    )
    l = lambdarepr(p)
    eval(h + l)
    assert l == "((((1) if (x &gt; 0) else (2))) if (y &gt; 0) else (3))"


def test_sum__1():
    # In each case, test eval() the lambdarepr() to make sure that
    # it evaluates to the same results as the symbolic expression
    s = Sum(x ** i, (i, a, b))
    l = lambdarepr(s)
    assert l == "(builtins.sum(x**i for i in range(a, b+1)))"

    args = x, a, b
    f = lambdify(args, s)
    v = 2, 3, 8
    assert f(*v) == s.subs(zip(args, v)).doit()

def test_sum__2():
    s = Sum(i * x, (i, a, b))
    l = lambdarepr(s)
    assert l == "(builtins.sum(i*x for i in range(a, b+1)))"

    args = x, a, b
    f = lambdify(args, s)
    v = 2, 3, 8
    assert f(*v) == s.subs(zip(args, v)).doit()


def test_multiple_sums():
    s = Sum(i * x + j, (i, a, b), (j, c, d))

    l = lambdarepr(s)
    assert l == "(builtins.sum(i*x + j for i in range(a, b+1) for j in range(c, d+1)))"

    args = x, a, b, c, d
    f = lambdify(args, s)
    vals = 2, 3, 4, 5, 6
    f_ref = s.subs(zip(args, vals)).doit()
    f_res = f(*vals)
    assert f_res == f_ref


def test_sqrt():
    prntr = LambdaPrinter({'standard' : 'python2'})
    assert prntr._print_Pow(sqrt(x), rational=False) == 'sqrt(x)'
    assert prntr._print_Pow(sqrt(x), rational=True) == 'x**(1./2.)'
    prntr = LambdaPrinter({'standard' : 'python3'})
    assert prntr._print_Pow(sqrt(x), rational=True) == 'x**(1/2)'


def test_settings():
    raises(TypeError, lambda: lambdarepr(sin(x), method="garbage"))


def test_numexpr():
    # test ITE rewrite as Piecewise
    from sympy.logic.boolalg import ITE
    expr = ITE(x &gt; 0, True, False, evaluate=False)
    assert NumExprPrinter().doprint(expr) == \
           "evaluate('where((x &gt; 0), True, False)', truediv=True)"


class CustomPrintedObject(Expr):
    def _lambdacode(self, printer):
        return 'lambda'

    def _tensorflowcode(self, printer):
        return 'tensorflow'

    def _numpycode(self, printer):
        return 'numpy'

    def _numexprcode(self, printer):
        return 'numexpr'

    def _mpmathcode(self, printer):
        return 'mpmath'


def test_printmethod():
    # In each case, printmethod is called to test
    # its working

    obj = CustomPrintedObject()
    assert LambdaPrinter().doprint(obj) == 'lambda'
    assert TensorflowPrinter().doprint(obj) == 'tensorflow'
    assert NumExprPrinter().doprint(obj) == "evaluate('numexpr', truediv=True)"

    assert NumExprPrinter().doprint(Piecewise((y, x &gt;= 0), (z, x &lt; 0))) == \
            "evaluate('where((x &gt;= 0), y, z)', truediv=True)"
</pre></body></html>