math2tex 1.0.0__py3-none-any.whl

math2tex/macros.py

@@ -0,0 +1,408 @@
+"""
+math2tex.macros — Macro definition, storage, and hygienic substitution.
+"""
+
+import copy
+import re
+import uuid
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Set, Tuple
+
+from .config import (
+    ALL_FUNC_NAMES, BUILTIN_LATEX_FUNCS, GREEK_LETTERS, SPECIAL_CONSTANTS,
+)
+from .ast_nodes import (
+    Abs, BigOp, BinOp, Derivative, Factorial, Fraction, Func,
+    InnerProduct, Integral, IntervalNode, MatrixNode, Node, Norm, Num,
+    PiecewiseNode, Power, Prime, SetBuilder, Subscript, TextLiteral,
+    UnOp, Var,
+)
+from .lexer import ParseError
+
+# =============================================================================
+# Macro Dataclass
+# =============================================================================
+
+@dataclass(frozen=True)
+class Macro:
+    name: str
+    args: List[str]
+    body: Node
+
+
+# =============================================================================
+# Maximum Expansion Depth
+# =============================================================================
+
+MAX_MACRO_DEPTH = 64
+
+# =============================================================================
+# Unique Symbol Generation
+# =============================================================================
+
+def _gensym(prefix: str = 'v') -> str:
+    return f'_m{prefix}_{uuid.uuid4().hex[:8]}'
+
+# =============================================================================
+# AST Inspection Helpers
+# =============================================================================
+
+def _iter_children(node: Node):
+    if isinstance(node, (Num, Var, TextLiteral)):
+        return
+    if isinstance(node, BinOp):
+        yield node.left; yield node.right
+    elif isinstance(node, UnOp):
+        yield node.expr
+    elif isinstance(node, Fraction):
+        yield node.num; yield node.den
+    elif isinstance(node, Power):
+        yield node.base; yield node.exp
+    elif isinstance(node, Subscript):
+        yield node.base; yield node.sub
+    elif isinstance(node, Func):
+        yield from node.args
+        if node.sub: yield node.sub
+        if node.power: yield node.power
+    elif isinstance(node, Factorial):
+        yield node.expr
+    elif isinstance(node, Prime):
+        yield node.expr
+    elif isinstance(node, Derivative):
+        yield node.expr
+        if node.order: yield node.order
+    elif isinstance(node, Integral):
+        yield node.expr
+        if node.lower: yield node.lower
+        if node.upper: yield node.upper
+    elif isinstance(node, BigOp):
+        if node.lower: yield node.lower
+        if node.upper: yield node.upper
+        yield node.expr
+    elif isinstance(node, Abs):
+        yield node.expr
+    elif isinstance(node, Norm):
+        yield node.expr
+    elif isinstance(node, InnerProduct):
+        yield node.left; yield node.right
+    elif isinstance(node, SetBuilder):
+        yield node.expr; yield node.condition
+    elif isinstance(node, MatrixNode):
+        for row in node.rows:
+            yield from row
+    elif isinstance(node, IntervalNode):
+        yield node.left_expr; yield node.right_expr
+    elif isinstance(node, PiecewiseNode):
+        for expr, cond in node.conditions:
+            yield expr
+            if cond: yield cond
+
+
+def free_vars(node: Node) -> Set[str]:
+    free: Set[str] = set()
+    bound: Set[str] = set()
+
+    def walk(n: Node, locally_bound: Set[str]) -> None:
+        if isinstance(n, Var):
+            if n.name and n.name not in locally_bound:
+                free.add(n.name)
+            return
+        if isinstance(n, (Num, TextLiteral)):
+            return
+
+        new_bound = set(locally_bound)
+        if isinstance(n, BigOp) and n.lower is not None:
+            bv = _extract_binding_name(n.lower)
+            if bv:
+                new_bound.add(bv)
+        if isinstance(n, Derivative):
+            new_bound.add(n.var)
+            if n.var2:
+                new_bound.add(n.var2)
+
+        for child in _iter_children(n):
+            walk(child, new_bound)
+
+    walk(node, set())
+    return free
+
+
+def _extract_binding_name(lower_node: Node) -> Optional[str]:
+    if isinstance(lower_node, BinOp) and lower_node.op in ('=', '->'):
+        if isinstance(lower_node.left, Var):
+            return lower_node.left.name
+    return None
+
+
+def find_bound_vars(node: Node) -> Set[str]:
+    names: Set[str] = set()
+
+    def walk(n: Node) -> None:
+        if isinstance(n, BigOp) and n.lower is not None:
+            bv = _extract_binding_name(n.lower)
+            if bv:
+                names.add(bv)
+        if isinstance(n, Derivative):
+            names.add(n.var)
+            if n.var2:
+                names.add(n.var2)
+        for child in _iter_children(n):
+            walk(child)
+
+    walk(node)
+    return names
+
+# =============================================================================
+# Alpha-Renaming
+# =============================================================================
+
+def _rename_var(node: Node, old_name: str, new_name: str) -> Node:
+    if isinstance(node, Var):
+        return Var(new_name) if node.name == old_name else node
+    if isinstance(node, Num):
+        return node
+    if isinstance(node, TextLiteral):
+        return node
+    if isinstance(node, BinOp):
+        return BinOp(node.op,
+                     _rename_var(node.left, old_name, new_name),
+                     _rename_var(node.right, old_name, new_name))
+    if isinstance(node, UnOp):
+        return UnOp(node.op, _rename_var(node.expr, old_name, new_name))
+    if isinstance(node, Fraction):
+        return Fraction(_rename_var(node.num, old_name, new_name),
+                        _rename_var(node.den, old_name, new_name))
+    if isinstance(node, Power):
+        return Power(_rename_var(node.base, old_name, new_name),
+                     _rename_var(node.exp, old_name, new_name))
+    if isinstance(node, Subscript):
+        return Subscript(_rename_var(node.base, old_name, new_name),
+                         _rename_var(node.sub, old_name, new_name))
+    if isinstance(node, Func):
+        return Func(
+            node.name,
+            [_rename_var(a, old_name, new_name) for a in node.args],
+            _rename_var(node.sub, old_name, new_name) if node.sub else None,
+            _rename_var(node.power, old_name, new_name) if node.power else None,
+            node.primes,
+        )
+    if isinstance(node, Factorial):
+        return Factorial(_rename_var(node.expr, old_name, new_name))
+    if isinstance(node, Prime):
+        return Prime(_rename_var(node.expr, old_name, new_name), node.order)
+    if isinstance(node, Derivative):
+        var = new_name if node.var == old_name else node.var
+        var2 = new_name if node.var2 == old_name else node.var2
+        return Derivative(
+            _rename_var(node.expr, old_name, new_name),
+            var,
+            _rename_var(node.order, old_name, new_name) if node.order else None,
+            node.is_partial,
+            var2,
+        )
+    if isinstance(node, Integral):
+        return Integral(
+            _rename_var(node.expr, old_name, new_name),
+            _rename_var(node.lower, old_name, new_name) if node.lower else None,
+            _rename_var(node.upper, old_name, new_name) if node.upper else None,
+            node.kind,
+        )
+    if isinstance(node, BigOp):
+        return BigOp(
+            node.op_name,
+            _rename_var(node.lower, old_name, new_name) if node.lower else None,
+            _rename_var(node.upper, old_name, new_name) if node.upper else None,
+            _rename_var(node.expr, old_name, new_name),
+        )
+    if isinstance(node, Abs):
+        return Abs(_rename_var(node.expr, old_name, new_name))
+    if isinstance(node, Norm):
+        return Norm(_rename_var(node.expr, old_name, new_name))
+    if isinstance(node, InnerProduct):
+        return InnerProduct(_rename_var(node.left, old_name, new_name),
+                            _rename_var(node.right, old_name, new_name))
+    if isinstance(node, SetBuilder):
+        return SetBuilder(_rename_var(node.expr, old_name, new_name),
+                          _rename_var(node.condition, old_name, new_name))
+    if isinstance(node, MatrixNode):
+        return MatrixNode(
+            [[_rename_var(c, old_name, new_name) for c in row]
+             for row in node.rows],
+            node.kind,
+        )
+    if isinstance(node, IntervalNode):
+        return IntervalNode(
+            node.left_delim,
+            _rename_var(node.left_expr, old_name, new_name),
+            _rename_var(node.right_expr, old_name, new_name),
+            node.right_delim,
+        )
+    if isinstance(node, PiecewiseNode):
+        return PiecewiseNode([
+            (_rename_var(e, old_name, new_name),
+             _rename_var(c, old_name, new_name) if c else None)
+            for e, c in node.conditions
+        ])
+
+    raise ValueError(f"Cannot rename in node of type {type(node)}")
+
+# =============================================================================
+# Clean Substitution
+# =============================================================================
+
+def substitute(node: Node, bindings: Dict[str, Node],
+               _depth: int = 0) -> Node:
+    if _depth > MAX_MACRO_DEPTH:
+        raise ParseError(
+            f"Maximum macro expansion depth ({MAX_MACRO_DEPTH}) exceeded "
+            f"— possible infinite recursion in macro definitions."
+        )
+
+    if not bindings:
+        return copy.deepcopy(node)
+
+    arg_free: Set[str] = set()
+    for arg_node in bindings.values():
+        arg_free |= free_vars(arg_node)
+
+    body_bound = find_bound_vars(node)
+
+    collisions = body_bound & (set(bindings.keys()) | arg_free)
+    renamed_node = node
+    for bv in collisions:
+        new_name = _gensym(bv)
+        renamed_node = _rename_var(renamed_node, bv, new_name)
+
+    return _do_substitute(renamed_node, bindings, _depth)
+
+
+def _do_substitute(node: Node, bindings: Dict[str, Node],
+                   _depth: int) -> Node:
+    if isinstance(node, Var):
+        if node.name in bindings:
+            return copy.deepcopy(bindings[node.name])
+        return node
+    if isinstance(node, Num):
+        return node
+    if isinstance(node, TextLiteral):
+        return node
+    if isinstance(node, BinOp):
+        return BinOp(node.op,
+                     _do_substitute(node.left, bindings, _depth),
+                     _do_substitute(node.right, bindings, _depth))
+    if isinstance(node, UnOp):
+        return UnOp(node.op, _do_substitute(node.expr, bindings, _depth))
+    if isinstance(node, Fraction):
+        return Fraction(_do_substitute(node.num, bindings, _depth),
+                        _do_substitute(node.den, bindings, _depth))
+    if isinstance(node, Power):
+        return Power(_do_substitute(node.base, bindings, _depth),
+                     _do_substitute(node.exp, bindings, _depth))
+    if isinstance(node, Subscript):
+        return Subscript(_do_substitute(node.base, bindings, _depth),
+                         _do_substitute(node.sub, bindings, _depth))
+    if isinstance(node, Func):
+        return Func(
+            node.name,
+            [_do_substitute(a, bindings, _depth) for a in node.args],
+            _do_substitute(node.sub, bindings, _depth) if node.sub else None,
+            _do_substitute(node.power, bindings, _depth) if node.power else None,
+            node.primes,
+        )
+    if isinstance(node, Factorial):
+        return Factorial(_do_substitute(node.expr, bindings, _depth))
+    if isinstance(node, Prime):
+        return Prime(_do_substitute(node.expr, bindings, _depth), node.order)
+    if isinstance(node, Derivative):
+        return Derivative(
+            _do_substitute(node.expr, bindings, _depth),
+            node.var,
+            _do_substitute(node.order, bindings, _depth) if node.order else None,
+            node.is_partial,
+            node.var2,
+        )
+    if isinstance(node, Integral):
+        return Integral(
+            _do_substitute(node.expr, bindings, _depth),
+            _do_substitute(node.lower, bindings, _depth) if node.lower else None,
+            _do_substitute(node.upper, bindings, _depth) if node.upper else None,
+            node.kind,
+        )
+    if isinstance(node, BigOp):
+        return BigOp(
+            node.op_name,
+            _do_substitute(node.lower, bindings, _depth) if node.lower else None,
+            _do_substitute(node.upper, bindings, _depth) if node.upper else None,
+            _do_substitute(node.expr, bindings, _depth),
+        )
+    if isinstance(node, Abs):
+        return Abs(_do_substitute(node.expr, bindings, _depth))
+    if isinstance(node, Norm):
+        return Norm(_do_substitute(node.expr, bindings, _depth))
+    if isinstance(node, InnerProduct):
+        return InnerProduct(_do_substitute(node.left, bindings, _depth),
+                            _do_substitute(node.right, bindings, _depth))
+    if isinstance(node, SetBuilder):
+        return SetBuilder(_do_substitute(node.expr, bindings, _depth),
+                          _do_substitute(node.condition, bindings, _depth))
+    if isinstance(node, MatrixNode):
+        return MatrixNode(
+            [[_do_substitute(c, bindings, _depth) for c in row]
+             for row in node.rows],
+            node.kind,
+        )
+    if isinstance(node, IntervalNode):
+        return IntervalNode(
+            node.left_delim,
+            _do_substitute(node.left_expr, bindings, _depth),
+            _do_substitute(node.right_expr, bindings, _depth),
+            node.right_delim,
+        )
+    if isinstance(node, PiecewiseNode):
+        return PiecewiseNode([
+            (_do_substitute(e, bindings, _depth),
+             _do_substitute(c, bindings, _depth) if c else None)
+            for e, c in node.conditions
+        ])
+
+    raise ValueError(f"Cannot substitute node of type {type(node)}")
+
+# =============================================================================
+# Macro Definition Parser
+# =============================================================================
+
+def parse_macro_def(line: str) -> Tuple[str, Macro]:
+    from .parser import Parser
+    from .lexer import _preprocess, tokenize
+
+    line = line.strip()[len("define "):].strip()
+    if '=' not in line:
+        raise ValueError("Macro definition must contain '='")
+    left, right = line.split('=', 1)
+    left, right = left.strip(), right.strip()
+
+    match = re.match(r'^([a-zA-Z_][a-zA-Z0-9_]*)(?:\((.*?)\))?$', left)
+    if not match:
+        raise ValueError("Invalid macro signature.")
+
+    name = match.group(1)
+    if (name in ALL_FUNC_NAMES or name in GREEK_LETTERS
+            or name in SPECIAL_CONSTANTS or name in BUILTIN_LATEX_FUNCS):
+        raise ValueError(
+            f"Cannot overwrite built-in function or constant '{name}'."
+        )
+
+    args_str = match.group(2)
+    args: List[str] = []
+    if args_str:
+        args = [a.strip() for a in args_str.split(',')]
+        for a in args:
+            if not re.match(r'^[a-zA-Z]+$', a):
+                raise ValueError(f"Invalid argument name '{a}' in macro.")
+
+    preprocessed, text_map = _preprocess(right)
+    tokens = tokenize(preprocessed)
+    parser = Parser(tokens, text_map=text_map)
+    body_ast = parser.parse()
+
+    return name, Macro(name, args, body_ast)