pythonocc-step-editor/step_editor/relation_formulas.py

from __future__ import annotations

from dataclasses import dataclass
import ast
import math
import re
from typing import Callable, Iterable


RELATION_REF_PATTERN = re.compile(
    r"\b(?P<kind>Face|Edge)(?P<object_id>\d+)\.(?P<parameter>[A-Za-z0-9_\u4e00-\u9fff]+)\b"
)


class RelationFormulaError(ValueError):
    pass


@dataclass(frozen=True)
class ObjectParameterRef:
    kind: str
    object_id: int
    parameter: str

    @property
    def token(self) -> str:
        return f"{self.kind}{self.object_id}.{self.parameter}"


@dataclass(frozen=True)
class RelationFormula:
    text: str
    target: ObjectParameterRef
    expression: str
    safe_expression: str
    references: tuple[ObjectParameterRef, ...]


class Vector3:
    __slots__ = ("values",)

    def __init__(self, values: Iterable[object]) -> None:
        items = tuple(values)
        if len(items) != 3:
            raise TypeError("Vector expression must contain exactly 3 values.")
        try:
            self.values = (float(items[0]), float(items[1]), float(items[2]))
        except (TypeError, ValueError) as exc:
            raise TypeError("Vector expression values must be numbers.") from exc

    def __iter__(self):
        return iter(self.values)

    def __len__(self) -> int:
        return 3

    def __getitem__(self, index: int) -> float:
        return self.values[index]

    def __repr__(self) -> str:
        return f"Vector3({self.values!r})"

    def __add__(self, other: object) -> "Vector3":
        right = _coerce_vector(other)
        return Vector3((self.values[0] + right[0], self.values[1] + right[1], self.values[2] + right[2]))

    def __radd__(self, other: object) -> "Vector3":
        return self.__add__(other)

    def __sub__(self, other: object) -> "Vector3":
        right = _coerce_vector(other)
        return Vector3((self.values[0] - right[0], self.values[1] - right[1], self.values[2] - right[2]))

    def __rsub__(self, other: object) -> "Vector3":
        left = _coerce_vector(other)
        return Vector3((left[0] - self.values[0], left[1] - self.values[1], left[2] - self.values[2]))

    def __mul__(self, other: object) -> "Vector3":
        scalar = _coerce_number(other)
        return Vector3((self.values[0] * scalar, self.values[1] * scalar, self.values[2] * scalar))

    def __rmul__(self, other: object) -> "Vector3":
        return self.__mul__(other)

    def __truediv__(self, other: object) -> "Vector3":
        scalar = _coerce_number(other)
        if abs(scalar) <= 1e-15:
            raise ZeroDivisionError("Vector division by zero.")
        return Vector3((self.values[0] / scalar, self.values[1] / scalar, self.values[2] / scalar))

    def __neg__(self) -> "Vector3":
        return Vector3((-self.values[0], -self.values[1], -self.values[2]))


def parse_relation_formula(text: str) -> RelationFormula:
    normalized = " ".join(str(text or "").strip().split())
    if not normalized:
        raise RelationFormulaError("请输入关系式。")
    if normalized.count("=") != 1:
        raise RelationFormulaError("关系式必须且只能包含一个等号，例如 Face87.直径 = Face85.直径。")
    left, expression = (part.strip() for part in normalized.split("=", 1))
    if not left or not expression:
        raise RelationFormulaError("关系式左侧和右侧都不能为空。")
    target_match = RELATION_REF_PATTERN.fullmatch(left)
    if target_match is None:
        raise RelationFormulaError("关系式左侧必须是 FaceID.参数 或 EdgeID.参数，例如 Face87.直径。")
    target = _ref_from_match(target_match)

    references: list[ObjectParameterRef] = []

    def replace_ref(match: re.Match[str]) -> str:
        references.append(_ref_from_match(match))
        return f"__ref{len(references) - 1}"

    safe_expression = RELATION_REF_PATTERN.sub(replace_ref, expression)
    try:
        tree = ast.parse(safe_expression, mode="eval")
    except SyntaxError as exc:
        raise RelationFormulaError(f"关系式右侧语法错误：{exc.msg}") from exc
    _validate_expression_tree(tree, len(references))
    return RelationFormula(
        text=f"{target.token} = {expression}",
        target=target,
        expression=expression,
        safe_expression=safe_expression,
        references=tuple(references),
    )


def evaluate_relation_formula(
    formula: RelationFormula,
    value_resolver: Callable[[ObjectParameterRef], object],
) -> float | Vector3:
    namespace: dict[str, object] = {}
    for index, ref in enumerate(formula.references):
        namespace[f"__ref{index}"] = _coerce_formula_value(value_resolver(ref))
    code = compile(formula.safe_expression, "<relation-formula>", "eval")
    try:
        value = eval(code, {"__builtins__": {}}, namespace)
    except ZeroDivisionError as exc:
        raise RelationFormulaError("关系式中出现除以 0。") from exc
    except Exception as exc:
        raise RelationFormulaError(f"关系式计算失败：{exc}") from exc
    return _coerce_formula_value(value)


def validate_relation_formula_graph(formulas: Iterable[RelationFormula]) -> None:
    target_to_formula: dict[str, RelationFormula] = {}
    for formula in formulas:
        target_token = formula.target.token
        if target_token in target_to_formula:
            raise RelationFormulaError(f"同一目标参数只能由一条关系式控制：{target_token}。")
        target_to_formula[target_token] = formula

    target_tokens = set(target_to_formula)
    graph: dict[str, list[str]] = {}
    for target_token, formula in target_to_formula.items():
        reference_tokens = [ref.token for ref in formula.references]
        if target_token in reference_tokens:
            raise RelationFormulaError(f"关系式不能引用自身：{target_token}。")
        graph[target_token] = [ref_token for ref_token in reference_tokens if ref_token in target_tokens]

    visit_state: dict[str, str] = {}
    stack: list[str] = []

    def visit(token: str) -> None:
        state = visit_state.get(token)
        if state == "visiting":
            start_index = stack.index(token) if token in stack else 0
            cycle = [*stack[start_index:], token]
            raise RelationFormulaError(f"关系式存在循环依赖：{' -> '.join(cycle)}。")
        if state == "visited":
            return
        visit_state[token] = "visiting"
        stack.append(token)
        for dependency in graph.get(token, []):
            visit(dependency)
        stack.pop()
        visit_state[token] = "visited"

    for target_token in graph:
        visit(target_token)


def relation_value_to_text(value: object) -> str:
    value = _coerce_formula_value(value)
    if isinstance(value, Vector3):
        return ", ".join(_format_number(item) for item in value.values)
    return _format_number(float(value))


def rewrite_relation_formula_ids(text: str, face_id_map: dict[int, int], edge_id_map: dict[int, int] | None = None) -> str:
    edge_id_map = dict(edge_id_map or {})

    def replace(match: re.Match[str]) -> str:
        kind = str(match.group("kind"))
        object_id = int(match.group("object_id"))
        parameter = str(match.group("parameter"))
        if kind == "Face" and object_id in face_id_map:
            object_id = int(face_id_map[object_id])
        elif kind == "Edge" and object_id in edge_id_map:
            object_id = int(edge_id_map[object_id])
        return f"{kind}{object_id}.{parameter}"

    return RELATION_REF_PATTERN.sub(replace, text)


def _ref_from_match(match: re.Match[str]) -> ObjectParameterRef:
    return ObjectParameterRef(
        kind=str(match.group("kind")),
        object_id=int(match.group("object_id")),
        parameter=str(match.group("parameter")),
    )


def _validate_expression_tree(tree: ast.AST, ref_count: int) -> None:
    allowed = (
        ast.Expression,
        ast.BinOp,
        ast.UnaryOp,
        ast.Name,
        ast.Load,
        ast.Constant,
        ast.Tuple,
        ast.Add,
        ast.Sub,
        ast.Mult,
        ast.Div,
        ast.UAdd,
        ast.USub,
    )
    for node in ast.walk(tree):
        if not isinstance(node, allowed):
            raise RelationFormulaError("关系式只支持数字、对象参数、括号、向量和 + - * / 运算。")
        if isinstance(node, ast.Name):
            if not re.fullmatch(r"__ref\d+", node.id):
                raise RelationFormulaError(f"未知参数引用：{node.id}")
            index = int(node.id.replace("__ref", ""))
            if index < 0 or index >= ref_count:
                raise RelationFormulaError(f"未知参数引用：{node.id}")
        elif isinstance(node, ast.Constant):
            if not isinstance(node.value, (int, float)):
                raise RelationFormulaError("关系式常量只支持数字。")
            if isinstance(node.value, float) and not math.isfinite(node.value):
                raise RelationFormulaError("关系式数字不能是 NaN 或无穷大。")
        elif isinstance(node, ast.Tuple):
            if len(node.elts) != 3:
                raise RelationFormulaError("向量必须是 3 个数字，例如 (0, 0, -3.5)。")


def _coerce_formula_value(value: object) -> float | Vector3:
    if isinstance(value, Vector3):
        return value
    if isinstance(value, (tuple, list)):
        return Vector3(value)
    return _coerce_number(value)


def _coerce_vector(value: object) -> tuple[float, float, float]:
    if isinstance(value, Vector3):
        return value.values
    if isinstance(value, (tuple, list)):
        return Vector3(value).values
    raise TypeError("Vector operation requires another 3D vector.")


def _coerce_number(value: object) -> float:
    if isinstance(value, bool):
        raise TypeError("Boolean is not a valid numeric formula value.")
    if isinstance(value, (int, float)):
        number = float(value)
    else:
        raise TypeError(f"{value!r} is not a valid numeric formula value.")
    if not math.isfinite(number):
        raise TypeError("Formula value must be finite.")
    return number


def _format_number(value: float) -> str:
    if abs(value) < 5e-13:
        value = 0.0
    return f"{value:.12g}"