diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index f85ce5752..b594187aa 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -10,6 +10,8 @@ from OCP.TDocStd import TDocStd_Document from OCP.TDataStd import TDataStd_Name, TDataStd_TreeNode from OCP.STEPCAFControl import STEPCAFControl_Reader +from OCP.XSControl import XSControl_TransferReader +from OCP.StepRepr import StepRepr_RepresentationItem from OCP.XCAFDoc import ( XCAFDoc_ColorSurf, XCAFDoc_DocumentTool, @@ -165,7 +167,10 @@ def importStep(assy: AssemblyProtocol, path: str, unit: UnitLiterals = "MM"): # Transfer the contents of the STEP file to the document step_reader.Transfer(doc) - _importDoc(doc, assy) + # Collect the STEP model for possible subshape related workarounds + transfer_reader = step_reader.Reader().WS().TransferReader() + + _importDoc(doc, assy, transfer_reader) def importXbf(assy: AssemblyProtocol, path: str): @@ -228,7 +233,11 @@ def importXml(assy: AssemblyProtocol, path: str): _importDoc(doc, assy) -def _importDoc(doc: TDocStd_Document, assy: AssemblyProtocol): +def _importDoc( + doc: TDocStd_Document, + assy: AssemblyProtocol, + transfer_reader: XSControl_TransferReader | None = None, +): def _process_label(lbl: TDF_Label, parent: AssemblyProtocol): """ Recursive method to process the assembly in a top-down manner. @@ -348,6 +357,17 @@ def _process_label(lbl: TDF_Label, parent: AssemblyProtocol): # Save the color info via the assembly subshape mechanism current.addSubshape(Shape.cast(cur_shape), color=color) + # if provided, check the STEP data model to collect additional subshapes on solid level + if transfer_reader: + if isinstance( + entity := transfer_reader.EntityFromShapeResult( + final_shape, 1 + ), + StepRepr_RepresentationItem, + ): + if name := entity.Name().ToCString(): + current.addSubshape(cq_shape, name=name) + return parent # Shape and color tools for extracting XCAF data diff --git a/tests/test_assembly.py b/tests/test_assembly.py index d98514a81..b1b1960fa 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -2662,3 +2662,24 @@ def test_assembly_without_children_roundtrip(kind, tmpdir): loaded = cq.Assembly.load(path) assert loaded.toCompound().Volume() == approx(1) + + +@pytest.mark.parametrize("kind", ["step", "xml", "xbf"]) +def test_assembly_with_solid_subshape_roundtrip(kind, tmpdir): + """ + An assembly with a solid that is also tagged as a subshape. + """ + + assy = cq.Assembly(name="root").add(box(1, 1, 1), name="child") + + assy["child"].addSubshape(assy["child"].obj, name="SOLID_TAG") + assy["child"].addSubshape(assy["child"].obj.face(), name="FACE_TAG") + + path = os.path.join(tmpdir, f"solid_subshape.{kind}") + assy.export(path) + + loaded = cq.Assembly.load(path) + child_subshape_names = loaded["child"]._subshape_names.values() + + assert "FACE_TAG" in child_subshape_names + assert "SOLID_TAG" in child_subshape_names