assignment.py 1.61 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
# -*- coding: utf-8 -*-
2
import sympy as sp
3
4
from sympy.printing.latex import LatexPrinter

5
6
7
8
9
try:
    from sympy.codegen.ast import Assignment
except ImportError:
    Assignment = None

10
11
12
13
14
15
16
__all__ = ['Assignment']


def print_assignment_latex(printer, expr):
    """sympy cannot print Assignments as Latex. Thus, this function is added to the sympy Latex printer"""
    printed_lhs = printer.doprint(expr.lhs)
    printed_rhs = printer.doprint(expr.rhs)
17
    return "{printed_lhs} \leftarrow {printed_rhs}".format(printed_lhs=printed_lhs, printed_rhs=printed_rhs)
18
19


Martin Bauer's avatar
Martin Bauer committed
20
def assignment_str(assignment):
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
    return "{lhs} ← {rhs}".format(lhs=assignment.lhs, rhs=assignment.rhs)


if Assignment:

    Assignment.__str__ = assignment_str
    LatexPrinter._print_Assignment = print_assignment_latex

else:
    # back port for older sympy versions that don't have Assignment  yet

    class Assignment(sp.Rel):

        rel_op = ':='
        __slots__ = []
Martin Bauer's avatar
Martin Bauer committed
36

37
38
39
40
41
42
43
44
45
46
47
        def __new__(cls, lhs, rhs=0, **assumptions):
            from sympy.matrices.expressions.matexpr import (
                MatrixElement, MatrixSymbol)
            from sympy.tensor.indexed import Indexed
            lhs = sp.sympify(lhs)
            rhs = sp.sympify(rhs)
            # Tuple of things that can be on the lhs of an assignment
            assignable = (sp.Symbol, MatrixSymbol, MatrixElement, Indexed)
            if not isinstance(lhs, assignable):
                raise TypeError("Cannot assign to lhs of type %s." % type(lhs))
            return sp.Rel.__new__(cls, lhs, rhs, **assumptions)
Martin Bauer's avatar
Martin Bauer committed
48

49
50
        __str__ = assignment_str
        _print_Assignment = print_assignment_latex