pythonocc-step-editor/step_editor/export.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
from .step_io import _prepare_shape_for_step_export, _write_step


class ExportMixin:
    def export_all(self, filename: str | Path) -> None:
        export_shape = _compound_from_shapes(
            _prepare_shape_for_step_export(part.shape) for part in self.display_parts()
        )
        _write_step(export_shape, Path(filename))

    def export_quality_info(self, scope: str, target_id: int | None = None) -> dict[str, object]:
        if scope == "all":
            return _shape_quality_info("当前完整模型", self.shape, expect_solid=False)
        if scope == "part":
            if target_id is None:
                raise ValueError("Part id is required.")
            part = self.part_by_id(target_id)
            if part is None:
                raise ValueError(f"Unknown part id {target_id}")
            info = _shape_quality_info(f"零件 {part.id}: {part.name}", part.shape, expect_solid=True)
            info["part_id"] = part.id
            return info
        if scope == "solid":
            if target_id is None or target_id < 0 or target_id >= len(self.solids):
                raise ValueError(f"Unknown solid id {target_id}")
            part_id, solid = self.solids[target_id]
            info = _shape_quality_info(f"Solid {target_id}", solid, expect_solid=True)
            info["part_id"] = part_id
            info["solid_id"] = target_id
            return info
        if scope == "face":
            if target_id is None or target_id < 0 or target_id >= len(self.faces):
                raise ValueError(f"Unknown face id {target_id}")
            face_ids = self.face_region_ids(target_id)
            shape = self._face_region_shape(target_id)
            label = f"Face {target_id}" if len(face_ids) == 1 else f"Face region from face {target_id}"
            info = _shape_quality_info(label, shape, expect_solid=False)
            info["part_id"] = self.face_part_ids[target_id]
            info["solid_id"] = self.face_solid_ids[target_id]
            info["face_id"] = target_id
            info["face_region_ids"] = tuple(face_ids)
            info["face_region_count"] = len(face_ids)
            return info
        if scope == "edge":
            if target_id is None or target_id < 0 or target_id >= len(self.edges):
                raise ValueError(f"Unknown edge id {target_id}")
            info = _shape_quality_info(f"Edge {target_id}", self.edges[target_id], expect_solid=False)
            info["part_id"] = self.edge_part_ids[target_id]
            info["solid_id"] = self._edge_solid_id(target_id)
            info["edge_id"] = target_id
            return info
        if scope == "feature":
            if target_id is None or target_id < 0 or target_id >= len(self.faces):
                raise ValueError(f"Unknown feature source face id {target_id}")
            feature = self.feature_info(target_id)
            face_ids = _int_values(feature.get("feature_highlight_face_ids")) or [target_id]
            shape = _compound_from_shapes(self.faces[face_id] for face_id in face_ids if 0 <= face_id < len(self.faces))
            info = _shape_quality_info(f"Feature from face {target_id}", shape, expect_solid=False)
            info["part_id"] = self.face_part_ids[target_id]
            info["solid_id"] = self.face_solid_ids[target_id]
            info["face_id"] = target_id
            info["feature_face_ids"] = tuple(face_ids)
            return info
        raise ValueError(f"Unknown export quality scope: {scope}")

    def export_part(self, part_id: int, filename: str | Path) -> None:
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")
        _write_step(_prepare_shape_for_step_export(part.shape), Path(filename))

    def export_solid(self, solid_id: int, filename: str | Path) -> None:
        if solid_id < 0 or solid_id >= len(self.solids):
            raise ValueError(f"Unknown solid id {solid_id}")
        _write_step(_prepare_shape_for_step_export(self.solids[solid_id][1]), Path(filename))

    def export_face(self, face_id: int, filename: str | Path) -> None:
        if face_id < 0 or face_id >= len(self.faces):
            raise ValueError(f"Unknown face id {face_id}")
        _write_step(self._face_region_shape(face_id), Path(filename))

    def face_region_ids(self, face_id: int) -> list[int]:
        if face_id < 0 or face_id >= len(self.faces):
            raise ValueError(f"Unknown face id {face_id}")
        return self.connected_same_domain_face_ids(face_id) or [face_id]

    def face_region_boundary_edge_ids(self, face_id: int) -> list[int]:
        return self._region_boundary_edge_ids(self.face_region_ids(face_id))

    def _face_region_shape(self, face_id: int) -> TopoDS_Shape:
        face_ids = self.face_region_ids(face_id)
        shapes = [self.faces[item] for item in face_ids if 0 <= item < len(self.faces)]
        if not shapes:
            raise ValueError(f"Face region export did not find any valid faces for face {face_id}.")
        if len(shapes) == 1:
            return shapes[0]
        return _unify_same_domain_shape(_compound_from_shapes(shapes))

    def export_edge(self, edge_id: int, filename: str | Path) -> None:
        if edge_id < 0 or edge_id >= len(self.edges):
            raise ValueError(f"Unknown edge id {edge_id}")
        _write_step(self.edges[edge_id], Path(filename))

    def export_feature(self, face_id: int, filename: str | Path) -> None:
        if face_id < 0 or face_id >= len(self.faces):
            raise ValueError(f"Unknown feature source face id {face_id}")
        feature = self.feature_info(face_id)
        face_ids = _int_values(feature.get("feature_highlight_face_ids")) or [face_id]
        shapes = [self.faces[item] for item in face_ids if 0 <= item < len(self.faces)]
        if not shapes:
            raise ValueError("Feature export did not find any valid faces.")
        _write_step(_compound_from_shapes(shapes), Path(filename))