pythonocc-step-editor/step_editor/step_io.py

from __future__ import annotations

import re
from pathlib import Path

from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_Transform
from OCC.Core.BRep import BRep_Builder
from OCC.Core.BRepTools import breptools
from OCC.Core.IFSelect import IFSelect_RetDone
from OCC.Core.Interface import Interface_Static
from OCC.Core.STEPCAFControl import STEPCAFControl_Reader
from OCC.Core.STEPControl import STEPControl_AsIs, STEPControl_Reader, STEPControl_Writer
from OCC.Core.TDF import TDF_Label, TDF_LabelSequence
from OCC.Core.TDocStd import TDocStd_Document
from OCC.Core.TopLoc import TopLoc_Location
from OCC.Core.TopoDS import TopoDS_Shape
from OCC.Core.XCAFDoc import XCAFDoc_DocumentTool

from .geometry_utils import _compound_from_shapes, _repair_shape, _unify_same_domain_shape
from .model_types import PartNode


def _load_with_xcaf(path: Path, product_names: list[str]) -> tuple[list[PartNode], TopoDS_Shape]:
    doc = TDocStd_Document("pythonocc-step-document")
    shape_tool = XCAFDoc_DocumentTool.ShapeTool(doc.Main())

    reader = STEPCAFControl_Reader()
    reader.SetColorMode(True)
    reader.SetLayerMode(True)
    reader.SetNameMode(True)
    reader.SetMatMode(True)
    reader.SetGDTMode(True)
    status = reader.ReadFile(str(path))
    if status != IFSelect_RetDone:
        raise ValueError(f"Could not read STEP file: {path}")
    if not reader.Transfer(doc):
        raise ValueError(f"Could not transfer STEP document: {path}")

    parts: list[PartNode] = []
    free_shapes = TDF_LabelSequence()
    shape_tool.GetFreeShapes(free_shapes)

    def next_name(label: TDF_Label, index: int) -> str:
        label_name = str(label.GetLabelName()).strip()
        if label_name:
            return label_name
        if index - 1 < len(product_names):
            return product_names[index - 1]
        return f"Part {index}"

    def add_node(
        name: str,
        kind: str,
        shape: TopoDS_Shape,
        parent_id: int | None,
        depth: int,
        path_text: str,
    ) -> PartNode:
        node = PartNode(len(parts) + 1, name, kind, shape, parent_id, depth, path_text)
        parts.append(node)
        return node

    def transformed_shape(label: TDF_Label, locations: list[TopLoc_Location]) -> TopoDS_Shape:
        shape = shape_tool.GetShape(label)
        if shape.IsNull() or not locations:
            return shape
        location = TopLoc_Location()
        for loc in locations:
            location = location.Multiplied(loc)
        return BRepBuilderAPI_Transform(shape, location.Transformation()).Shape()

    def walk(label: TDF_Label, parent_id: int | None, depth: int, locations: list[TopLoc_Location], path_names: list[str]):
        name = next_name(label, len(parts) + 1)
        label_path = " / ".join(path_names + [name])

        if shape_tool.IsAssembly(label):
            node = add_node(name, "assembly", transformed_shape(label, locations), parent_id, depth, label_path)
            components = TDF_LabelSequence()
            shape_tool.GetComponents(label, components)
            for i in range(1, components.Length() + 1):
                component = components.Value(i)
                if shape_tool.IsReference(component):
                    referred = TDF_Label()
                    shape_tool.GetReferredShape(component, referred)
                    loc = shape_tool.GetLocation(component)
                    walk(referred, node.id, depth + 1, locations + [loc], path_names + [name])
                else:
                    walk(component, node.id, depth + 1, locations, path_names + [name])
            return

        if shape_tool.IsSimpleShape(label) or shape_tool.IsShape(label):
            add_node(name, "part", transformed_shape(label, locations), parent_id, depth, label_path)

    for i in range(1, free_shapes.Length() + 1):
        walk(free_shapes.Value(i), None, 0, [], [])

    display_shapes = [p.shape for p in parts if p.kind == "part" and not p.shape.IsNull()]
    if not display_shapes:
        display_shapes = [p.shape for p in parts if not p.shape.IsNull()]
    return parts, _compound_from_shapes(display_shapes)


def _load_plain_step(path: Path) -> TopoDS_Shape:
    reader = STEPControl_Reader()
    status = reader.ReadFile(str(path))
    if status != IFSelect_RetDone:
        raise ValueError(f"Could not read STEP file: {path}")
    if not reader.TransferRoots():
        raise ValueError(f"Could not transfer STEP roots: {path}")
    return reader.Shape()


def _parse_product_names(path: Path) -> list[str]:
    text = path.read_text(errors="ignore")
    names = re.findall(r"PRODUCT\('((?:''|[^'])*)'", text)
    return [name.replace("''", "'") for name in names if name.strip()]


def _write_step(shape: TopoDS_Shape, filename: Path) -> None:
    if shape.IsNull():
        raise ValueError("Cannot export a null shape.")
    filename.parent.mkdir(parents=True, exist_ok=True)
    Interface_Static.SetCVal("write.step.schema", "AP214IS")
    writer = STEPControl_Writer()
    writer.Transfer(shape, STEPControl_AsIs)
    status = writer.Write(str(filename))
    if status != IFSelect_RetDone:
        raise IOError(f"Could not write STEP file: {filename}")


def _load_brep_shape(filename: str | Path) -> TopoDS_Shape:
    path = Path(filename)
    shape = TopoDS_Shape()
    builder = BRep_Builder()
    if not breptools.Read(shape, str(path), builder) or shape.IsNull():
        raise IOError(f"Could not read BREP file: {path}")
    return shape


def _write_brep(shape: TopoDS_Shape, filename: str | Path) -> None:
    if shape.IsNull():
        raise ValueError("Cannot export a null shape.")
    path = Path(filename)
    path.parent.mkdir(parents=True, exist_ok=True)
    if not breptools.Write(shape, str(path)):
        raise IOError(f"Could not write BREP file: {path}")


def _prepare_shape_for_step_export(shape: TopoDS_Shape) -> TopoDS_Shape:
    if shape.IsNull():
        return shape
    try:
        repaired = _repair_shape(shape)
        unified = _unify_same_domain_shape(repaired)
        repaired_unified = _repair_shape(unified)
        if repaired_unified.IsNull():
            return shape
        return repaired_unified
    except Exception:
        return shape