expression_lib_example.py

from adt import Case, adt


@adt
class Expression:
    LITERAL: Case[float]
    ADD: Case["Expression", "Expression"]
    SUBTRACT: Case["Expression", "Expression"]
    MULTIPLY: Case["Expression", "Expression"]
    DIVIDE: Case["Expression", "Expression"]


def evaluate(expr: Expression):
    return expr.match(
        literal=lambda l: l,
        add=lambda x, y: evaluate(x) + evaluate(y),
        subtract=lambda x, y: evaluate(x) - evaluate(y),
        multiply=lambda x, y: evaluate(x) * evaluate(y),
        divide=lambda x, y: evaluate(x) / evaluate(y),
    )


def test_expression_evaluation():
    expression = Expression.MULTIPLY(
        Expression.ADD(Expression.LITERAL(1.0), Expression.LITERAL(2.0)),
        Expression.LITERAL(5.0),
    )
    value = evaluate(expression)
    assert value == (1 + 2) * 5