pythonocc-step-editor/step_editor/transforms.py

from __future__ import annotations

import math
from pathlib import Path
from typing import Callable, Iterable

from OCC.Core.BRep import BRep_Tool
from OCC.Core.BRepAdaptor import BRepAdaptor_Curve, BRepAdaptor_Surface
from OCC.Core.BRepAlgoAPI import BRepAlgoAPI_Cut, BRepAlgoAPI_Defeaturing, BRepAlgoAPI_Fuse
from OCC.Core.BRepBndLib import brepbndlib
from OCC.Core.BOPAlgo import BOPAlgo_GlueFull
from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_Transform
from OCC.Core.BRepCheck import BRepCheck_Analyzer
from OCC.Core.BRepClass3d import BRepClass3d_SolidClassifier
from OCC.Core.BRepFilletAPI import BRepFilletAPI_MakeChamfer, BRepFilletAPI_MakeFillet
from OCC.Core.BRepGProp import brepgprop
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
from OCC.Core.BRepPrimAPI import BRepPrimAPI_MakeCylinder, BRepPrimAPI_MakePrism
from OCC.Core.Bnd import Bnd_Box
from OCC.Core.GeomAbs import (
    GeomAbs_BSplineCurve,
    GeomAbs_BSplineSurface,
    GeomAbs_BezierCurve,
    GeomAbs_BezierSurface,
    GeomAbs_Circle,
    GeomAbs_Cone,
    GeomAbs_Cylinder,
    GeomAbs_Ellipse,
    GeomAbs_Hyperbola,
    GeomAbs_Line,
    GeomAbs_OffsetSurface,
    GeomAbs_OtherCurve,
    GeomAbs_OtherSurface,
    GeomAbs_Parabola,
    GeomAbs_Plane,
    GeomAbs_Sphere,
    GeomAbs_SurfaceOfExtrusion,
    GeomAbs_SurfaceOfRevolution,
    GeomAbs_Torus,
)
from OCC.Core.GProp import GProp_GProps
from OCC.Core.ShapeFix import ShapeFix_Shape
from OCC.Core.ShapeUpgrade import ShapeUpgrade_UnifySameDomain
from OCC.Core.TopAbs import (
    TopAbs_EDGE,
    TopAbs_EXTERNAL,
    TopAbs_FACE,
    TopAbs_FORWARD,
    TopAbs_IN,
    TopAbs_INTERNAL,
    TopAbs_OUT,
    TopAbs_REVERSED,
    TopAbs_SOLID,
)
from OCC.Core.TopExp import TopExp_Explorer, topexp
from OCC.Core.TopLoc import TopLoc_Location
from OCC.Core.TopoDS import TopoDS_Compound, TopoDS_Shape, topods
from OCC.Core.TopTools import TopTools_IndexedDataMapOfShapeListOfShape, TopTools_IndexedMapOfShape
from OCC.Core.gp import gp_Ax1, gp_Ax2, gp_Dir, gp_Pnt, gp_Trsf, gp_Vec
from OCC.Extend.TopologyUtils import TopologyExplorer, discretize_edge

from .constants import CURVE_TYPES, SNAPSHOT_FACE_LOGICAL_IDS_KEY, SURFACE_TYPES
from .geometry_utils import *  # noqa: F403


class TransformMixin:
    def translate_part_plan(self, part_id: int, vector: tuple[float, float, float]) -> dict[str, object]:
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")
        readiness = _translation_readiness(vector, part.shape)
        return {
            "status": readiness["translate_status"],
            "risk": readiness["translate_risk"],
            "message": readiness["translate_note"],
            "warnings": readiness["translate_warnings"],
            "blockers": readiness["translate_blockers"],
            "target_kind": "part",
            "part_id": part.id,
            "name": part.name,
            "translation_vector": vector,
            "translation_distance": _vector_length(vector),
            "bbox_diagonal": _shape_diagonal(part.shape),
        }

    def translate_solid_plan(self, solid_id: int, vector: tuple[float, float, float]) -> dict[str, object]:
        if solid_id < 0 or solid_id >= len(self.solids):
            raise ValueError(f"Unknown solid id {solid_id}")
        part_id, solid = self.solids[solid_id]
        readiness = _translation_readiness(vector, solid)
        part = self.part_by_id(part_id)
        part_solid_count = len(_explore(part.shape, TopAbs_SOLID)) if part is not None else 0
        warnings = readiness["translate_warnings"]
        risk = readiness["translate_risk"]
        status = readiness["translate_status"]
        if part_solid_count <= 1 and status != "blocked":
            warnings = _join_nonempty(warnings, "当前零件只有一个Solid，平移Solid实际会移动整个零件 shape。")
            if risk == "low":
                risk = "medium"
                status = "caution"
        return {
            "status": status,
            "risk": risk,
            "message": _join_nonempty(readiness["translate_note"], warnings),
            "warnings": warnings,
            "blockers": readiness["translate_blockers"],
            "target_kind": "solid",
            "part_id": part_id,
            "solid_id": solid_id,
            "part_solid_count": part_solid_count,
            "translation_vector": vector,
            "translation_distance": _vector_length(vector),
            "bbox_diagonal": _shape_diagonal(solid),
        }

    def translate_part(self, part_id: int, vector: tuple[float, float, float]) -> str:
        plan = self.translate_part_plan(part_id, vector)
        if plan["status"] == "blocked":
            raise ValueError(str(plan["message"]))
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"未知零件 ID {part_id}")
        part.shape = _translated_shape_by_vector(part.shape, vector)
        _ensure_valid_shape(part.shape)
        self.refresh_topology()
        return (
            f"零件已平移: 零件 {part_id}, vector={_format_tuple(vector)}, "
            f"distance={float(plan['translation_distance']):g}, risk={plan['risk']}."
        )

    def translate_solid(self, solid_id: int, vector: tuple[float, float, float]) -> str:
        plan = self.translate_solid_plan(solid_id, vector)
        if plan["status"] == "blocked":
            raise ValueError(str(plan["message"]))
        part_id, solid = self.solids[solid_id]
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")

        part_solids = _explore(part.shape, TopAbs_SOLID)
        if len(part_solids) <= 1:
            part.shape = _translated_shape_by_vector(part.shape, vector)
        else:
            translated = _translated_shape_by_vector(solid, vector)
            replaced = False
            shapes: list[TopoDS_Shape] = []
            for item in part_solids:
                if not replaced and _same_shape(item, solid):
                    shapes.append(translated)
                    replaced = True
                else:
                    shapes.append(item)
            if not replaced:
                raise RuntimeError(f"Could not locate solid {solid_id} inside part {part_id}.")
            part.shape = _compound_from_shapes(shapes)

        _ensure_valid_shape(part.shape)
        self.refresh_topology()
        return (
            f"Solid translated: solid {solid_id}, part {part_id}, vector={_format_tuple(vector)}, "
            f"distance={float(plan['translation_distance']):g}, risk={plan['risk']}."
        )

    def rotate_part_plan(self, part_id: int, axis: str, angle_degrees: float) -> dict[str, object]:
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")
        readiness = _rotation_readiness(axis, angle_degrees)
        return {
            "status": readiness["rotate_status"],
            "risk": readiness["rotate_risk"],
            "message": readiness["rotate_note"],
            "warnings": readiness["rotate_warnings"],
            "blockers": readiness["rotate_blockers"],
            "target_kind": "part",
            "part_id": part.id,
            "name": part.name,
            "rotation_axis": axis.upper(),
            "rotation_angle_degrees": angle_degrees,
            "rotation_center": _shape_center(part.shape),
            "bbox_diagonal": _shape_diagonal(part.shape),
        }

    def rotate_solid_plan(self, solid_id: int, axis: str, angle_degrees: float) -> dict[str, object]:
        if solid_id < 0 or solid_id >= len(self.solids):
            raise ValueError(f"Unknown solid id {solid_id}")
        part_id, solid = self.solids[solid_id]
        readiness = _rotation_readiness(axis, angle_degrees)
        part = self.part_by_id(part_id)
        part_solid_count = len(_explore(part.shape, TopAbs_SOLID)) if part is not None else 0
        warnings = readiness["rotate_warnings"]
        risk = readiness["rotate_risk"]
        status = readiness["rotate_status"]
        if part_solid_count <= 1 and status != "blocked":
            warnings = _join_nonempty(warnings, "当前零件只有一个Solid，旋转Solid实际会旋转整个零件 shape。")
            if risk == "low":
                risk = "medium"
                status = "caution"
        return {
            "status": status,
            "risk": risk,
            "message": _join_nonempty(readiness["rotate_note"], warnings),
            "warnings": warnings,
            "blockers": readiness["rotate_blockers"],
            "target_kind": "solid",
            "part_id": part_id,
            "solid_id": solid_id,
            "part_solid_count": part_solid_count,
            "rotation_axis": axis.upper(),
            "rotation_angle_degrees": angle_degrees,
            "rotation_center": _shape_center(solid),
            "bbox_diagonal": _shape_diagonal(solid),
        }

    def rotate_part(self, part_id: int, axis: str, angle_degrees: float) -> str:
        plan = self.rotate_part_plan(part_id, axis, angle_degrees)
        if plan["status"] == "blocked":
            raise ValueError(str(plan["message"]))
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"未知零件 ID {part_id}")
        part.shape = _rotated_shape(part.shape, str(plan["rotation_axis"]), float(plan["rotation_angle_degrees"]), plan["rotation_center"])
        _ensure_valid_shape(part.shape)
        self.refresh_topology()
        return (
            f"零件已旋转: 零件 {part_id}, axis={plan['rotation_axis']}, "
            f"angle={float(plan['rotation_angle_degrees']):g}, center={_format_tuple(plan['rotation_center'])}, "
            f"risk={plan['risk']}."
        )

    def rotate_solid(self, solid_id: int, axis: str, angle_degrees: float) -> str:
        plan = self.rotate_solid_plan(solid_id, axis, angle_degrees)
        if plan["status"] == "blocked":
            raise ValueError(str(plan["message"]))
        part_id, solid = self.solids[solid_id]
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")

        part_solids = _explore(part.shape, TopAbs_SOLID)
        if len(part_solids) <= 1:
            part.shape = _rotated_shape(part.shape, str(plan["rotation_axis"]), float(plan["rotation_angle_degrees"]), plan["rotation_center"])
        else:
            rotated = _rotated_shape(solid, str(plan["rotation_axis"]), float(plan["rotation_angle_degrees"]), plan["rotation_center"])
            replaced = False
            shapes: list[TopoDS_Shape] = []
            for item in part_solids:
                if not replaced and _same_shape(item, solid):
                    shapes.append(rotated)
                    replaced = True
                else:
                    shapes.append(item)
            if not replaced:
                raise RuntimeError(f"Could not locate solid {solid_id} inside part {part_id}.")
            part.shape = _compound_from_shapes(shapes)

        _ensure_valid_shape(part.shape)
        self.refresh_topology()
        return (
            f"Solid rotated: solid {solid_id}, part {part_id}, axis={plan['rotation_axis']}, "
            f"angle={float(plan['rotation_angle_degrees']):g}, center={_format_tuple(plan['rotation_center'])}, "
            f"risk={plan['risk']}."
        )