File size: 4,866 Bytes
edede4c | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | """
Tests for the AST Generator (Layer 3).
"""
import pytest
import sympy as sp
from mathtok.ast_generator import ASTGenerator, ASTNode
from mathtok.canonicalizer import Canonicalizer
@pytest.fixture
def gen():
return ASTGenerator()
@pytest.fixture
def canon():
return Canonicalizer(do_simplify=False, do_expand=False)
def parse(expr_str: str):
from sympy.parsing.sympy_parser import (
parse_expr, standard_transformations,
implicit_multiplication_application, convert_xor,
)
return parse_expr(
expr_str,
transformations=standard_transformations + (
implicit_multiplication_application, convert_xor,
),
local_dict={"x": sp.Symbol("x"), "y": sp.Symbol("y"),
"a": sp.Symbol("a"), "b": sp.Symbol("b"),
"n": sp.Symbol("n")},
)
class TestBasicNodes:
def test_symbol(self, gen):
ast = gen.generate(sp.Symbol("x"))
assert ast.token == "VAR_X"
assert ast.is_leaf
def test_integer_zero(self, gen):
ast = gen.generate(sp.Integer(0))
assert ast.token == "CONST_0"
def test_integer_positive(self, gen):
ast = gen.generate(sp.Integer(5))
assert ast.token == "CONST_5"
def test_integer_negative(self, gen):
ast = gen.generate(sp.Integer(-3))
assert ast.token == "OP_NEG"
assert ast.children[0].token == "CONST_3"
def test_pi(self, gen):
ast = gen.generate(sp.pi)
assert ast.token == "CONST_PI"
def test_e(self, gen):
ast = gen.generate(sp.E)
assert ast.token == "CONST_E"
def test_rational(self, gen):
ast = gen.generate(sp.Rational(1, 2))
assert ast.token == "FRAC"
assert len(ast.children) == 2
class TestArithmetic:
def test_add(self, gen):
expr = parse("x + 1")
ast = gen.generate(expr)
assert ast.token == "OP_ADD"
tokens = gen.get_all_tokens(ast)
assert "VAR_X" in tokens
assert "CONST_1" in tokens
def test_mul(self, gen):
expr = parse("2*x")
ast = gen.generate(expr)
# 2*x is either OP_MUL or OP_NEG etc.
assert ast.token in ("OP_MUL", "VAR_X", "CONST_2")
def test_pow(self, gen):
expr = parse("x^2")
ast = gen.generate(expr)
assert ast.token == "OP_POW"
assert ast.children[0].token == "VAR_X"
assert ast.children[1].token == "CONST_2"
def test_negation(self, gen):
expr = sp.Mul(sp.Integer(-1), sp.Symbol("x"))
ast = gen.generate(expr)
assert ast.token == "OP_NEG"
def test_reciprocal(self, gen):
expr = sp.Pow(sp.Symbol("x"), sp.Integer(-1))
ast = gen.generate(expr)
assert ast.token == "OP_RECIP"
class TestFunctions:
def test_sin(self, gen):
expr = sp.sin(sp.Symbol("x"))
ast = gen.generate(expr)
assert ast.token == "FUNC_SIN"
assert ast.children[0].token == "VAR_X"
def test_cos(self, gen):
ast = gen.generate(sp.cos(sp.Symbol("x")))
assert ast.token == "FUNC_COS"
def test_exp(self, gen):
ast = gen.generate(sp.exp(sp.Symbol("x")))
assert ast.token == "FUNC_EXP"
def test_log(self, gen):
ast = gen.generate(sp.log(sp.Symbol("x")))
assert ast.token == "FUNC_LOG"
def test_sqrt(self, gen):
# SymPy represents sqrt(x) internally as Pow(x, Rational(1,2))
# so the AST correctly emits OP_POW; FUNC_SQRT is only emitted
# when sympy.sqrt is used directly before any canonicalization.
ast = gen.generate(sp.sqrt(sp.Symbol("x")))
# Accept either FUNC_SQRT (direct) or OP_POW (post-simplification)
assert ast.token in ("FUNC_SQRT", "OP_POW")
class TestTreeProperties:
def test_depth_assignment(self, gen):
expr = parse("x^2 + 1")
ast = gen.generate(expr)
assert ast.depth == 0
for child in ast.children:
assert child.depth == 1
def test_unique_node_ids(self, gen):
expr = parse("x^2 + 2*x + 1")
ast = gen.generate(expr)
all_ids: list[int] = []
def collect(node):
all_ids.append(node.node_id)
for c in node.children:
collect(c)
collect(ast)
assert len(all_ids) == len(set(all_ids)), "Node IDs must be unique"
def test_subtree_size(self, gen):
ast = gen.generate(sp.Integer(5))
assert ast.subtree_size == 1
expr = parse("x + 1")
ast = gen.generate(expr)
assert ast.subtree_size == 3 # ADD + VAR_X + CONST_1
def test_variable_extraction(self, gen):
expr = parse("x^2 + y + 1")
ast = gen.generate(expr)
vars_ = gen.get_variable_tokens(ast)
assert "VAR_X" in vars_
assert "VAR_Y" in vars_
|