From 821aa1e6de6d07a1a1ef067d48124258ec3c9dab Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Wed, 26 Feb 2025 11:32:11 -0500 Subject: [PATCH 01/44] Initial commit with colors and locations working for simple assemblies --- cadquery/assembly.py | 14 +++ cadquery/occ_impl/importers/assembly.py | 119 ++++++++++++++++++++++++ tests/test_assembly.py | 37 ++++++++ 3 files changed, 170 insertions(+) create mode 100644 cadquery/occ_impl/importers/assembly.py diff --git a/cadquery/assembly.py b/cadquery/assembly.py index bed1fdc8a..22d9ab068 100644 --- a/cadquery/assembly.py +++ b/cadquery/assembly.py @@ -34,6 +34,7 @@ exportGLTF, STEPExportModeLiterals, ) +from .occ_impl.importers.assembly import importStep as importStepTopLevel from .selectors import _expression_grammar as _selector_grammar from .utils import deprecate @@ -564,6 +565,19 @@ def export( return self + @staticmethod + def importStep(path: str) -> "Assembly": + """ + Reads an assembly from a STEP file. + + :param path: Path and filename for writing. + :return: An Assembly object. + """ + + assy = importStepTopLevel(path) + + return assy + @classmethod def load(cls, path: str) -> "Assembly": diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py new file mode 100644 index 000000000..410ece33c --- /dev/null +++ b/cadquery/occ_impl/importers/assembly.py @@ -0,0 +1,119 @@ +from OCP.TCollection import TCollection_ExtendedString +from OCP.Quantity import Quantity_Color, Quantity_ColorRGBA +from OCP.TDocStd import TDocStd_Document +from OCP.IFSelect import IFSelect_RetDone +from OCP.STEPCAFControl import STEPCAFControl_Reader +from OCP.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ColorGen, XCAFDoc_ColorSurf +from OCP.TDF import TDF_Label, TDF_LabelSequence + +import cadquery as cq +from ..assembly import AssemblyProtocol + + +def importStep(path: str) -> AssemblyProtocol: + """ + Import a step file into an assembly. + """ + + # The assembly that is being built from the step file + assy = cq.Assembly() + + # Document that the step file will be read into + doc = TDocStd_Document(TCollection_ExtendedString("XmlOcaf")) + + # Create and configure a STEP reader + step_reader = STEPCAFControl_Reader() + step_reader.SetColorMode(True) + step_reader.SetNameMode(True) + step_reader.SetLayerMode(True) + step_reader.SetSHUOMode(True) + + # Read the STEP file + status = step_reader.ReadFile(path) + if status != IFSelect_RetDone: + raise ValueError(f"Error reading STEP file: {path}") + + # Transfer the contents of the STEP file to the document + step_reader.Transfer(doc) + + # Shape and color tools for extracting XCAF data + shape_tool = XCAFDoc_DocumentTool.ShapeTool_s(doc.Main()) + color_tool = XCAFDoc_DocumentTool.ColorTool_s(doc.Main()) + + def process_label(label, parent_location=None): + """ + Recursive function that allows us to process the hierarchy of the assembly as represented + in the step file. + """ + + # Handle reference labels + if shape_tool.IsReference_s(label): + ref_label = TDF_Label() + shape_tool.GetReferredShape_s(label, ref_label) + process_label(ref_label, parent_location) + return + + # Process components + comp_labels = TDF_LabelSequence() + shape_tool.GetComponents_s(label, comp_labels) + for i in range(comp_labels.Length()): + sub_label = comp_labels.Value(i + 1) + + # The component level holds the location for its shapes + loc = shape_tool.GetLocation_s(sub_label) + if loc: + location = cq.Location(loc) + + # Make sure that the location object is actually doing something interesting + # This is done because the location may have to go through multiple levels of + # components before the shapes are found. This allows the top-level component + # to specify the location/rotation of the shapes. + if location.toTuple()[0] == (0, 0, 0) and location.toTuple()[1] == ( + 0, + 0, + 0, + ): + location = parent_location + else: + location = parent_location + + process_label(sub_label, location) + + # Check to see if we have an endpoint shape + if shape_tool.IsSimpleShape_s(label): + shape = shape_tool.GetShape_s(label) + + # Process the color for the shape, which could be of different types + color = Quantity_Color() + if color_tool.GetColor_s(label, XCAFDoc_ColorSurf, color): + r = color.Red() + g = color.Green() + b = color.Blue() + cq_color = cq.Color(r, g, b) + elif color_tool.GetColor_s(label, XCAFDoc_ColorGen, color): + r = color.Red() + g = color.Green() + b = color.Blue() + cq_color = cq.Color(r, g, b) + else: + cq_color = cq.Color(0.5, 0.5, 0.5) + + # Handle the location if it was passed down form a parent component + if parent_location is not None: + assy.add(cq.Shape.cast(shape), color=cq_color, loc=parent_location) + else: + assy.add(cq.Shape.cast(shape), color=cq_color) + + # Grab the labels, which should hold the assembly parent + labels = TDF_LabelSequence() + shape_tool.GetFreeShapes(labels) + + # Make sure that we are working with an assembly + if shape_tool.IsAssembly_s(labels.Value(1)): + # Start the recursive processing of the assembly + process_label(labels.Value(1)) + + else: + raise ValueError("Step file does not contain an assembly") + + return assy diff --git a/tests/test_assembly.py b/tests/test_assembly.py index d66da7467..4ba94dd3f 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -616,6 +616,43 @@ def test_step_export(nested_assy, tmp_path_factory): assert pytest.approx(c2.toTuple()) == (0, 4, 0) +def test_step_import(tmp_path_factory): + """ + Exports a simple assembly to step with locations and colors and ensures that information is + preserved when the assembly is re-imported from the step file. + """ + + # Use a temporary directory + tmpdir = tmp_path_factory.mktemp("out") + metadata_path = os.path.join(tmpdir, "metadata.step") + + # Assembly with all the metadata added that needs to be read by the STEP importer + cube_1 = cq.Workplane().box(10, 10, 10) + cube_2 = cq.Workplane().box(15, 15, 15) + assy = cq.Assembly() + assy.add(cube_1, name="cube_1", color=cq.Color(1.0, 0.0, 0.0)) + assy.add( + cube_2, + name="cube_2", + color=cq.Color(0.0, 1.0, 0.0), + loc=cq.Location((15, 15, 15)), + ) + exportAssembly(assy, metadata_path) + + assy = cq.Assembly.importStep(path=metadata_path) + + # Make sure that there are the correct number of parts in the assembly + assert len(assy.children) == 2 + + # Make sure that the colors are correct + assert assy.children[0].color == cq.Color(1.0, 0.0, 0.0) + assert assy.children[1].color == cq.Color(0.0, 1.0, 0.0) + + # Make sure that the locations are correct + assert assy.children[0].loc.toTuple()[0] == (0, 0, 0) + assert assy.children[1].loc.toTuple()[0] == (15, 15, 15) + + @pytest.mark.parametrize( "assy_fixture, expected", [ From c3d557329f06224546c81e858a0dd9ad606834f5 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Wed, 26 Feb 2025 11:56:17 -0500 Subject: [PATCH 02/44] Trying to fix mypy error --- cadquery/occ_impl/importers/assembly.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 410ece33c..be068d044 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -7,10 +7,9 @@ from OCP.TDF import TDF_Label, TDF_LabelSequence import cadquery as cq -from ..assembly import AssemblyProtocol -def importStep(path: str) -> AssemblyProtocol: +def importStep(path: str) -> "Assembly": """ Import a step file into an assembly. """ From da9844ac3fb77b2817fe7bd1b56f0dde0159fbbd Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Wed, 26 Feb 2025 15:42:17 -0500 Subject: [PATCH 03/44] Jumping through hoops to make mypy happy --- cadquery/assembly.py | 2 +- cadquery/occ_impl/importers/assembly.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cadquery/assembly.py b/cadquery/assembly.py index 22d9ab068..29ec18698 100644 --- a/cadquery/assembly.py +++ b/cadquery/assembly.py @@ -574,7 +574,7 @@ def importStep(path: str) -> "Assembly": :return: An Assembly object. """ - assy = importStepTopLevel(path) + assy = cast(Assembly, importStepTopLevel(path)) return assy diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index be068d044..b3cf64289 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -7,9 +7,9 @@ from OCP.TDF import TDF_Label, TDF_LabelSequence import cadquery as cq +from ..assembly import AssemblyProtocol - -def importStep(path: str) -> "Assembly": +def importStep(path: str) -> AssemblyProtocol: """ Import a step file into an assembly. """ From 7b674c85c270c92c866300b85798e045e1a5d486 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Wed, 26 Feb 2025 16:01:49 -0500 Subject: [PATCH 04/44] Lint fix --- cadquery/occ_impl/importers/assembly.py | 1 + 1 file changed, 1 insertion(+) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index b3cf64289..410ece33c 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -9,6 +9,7 @@ import cadquery as cq from ..assembly import AssemblyProtocol + def importStep(path: str) -> AssemblyProtocol: """ Import a step file into an assembly. From 288c0a001acd9c509f4768b5fa9211a781146547 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Thu, 27 Feb 2025 08:51:39 -0500 Subject: [PATCH 05/44] Refactored method interface based on suggestion --- cadquery/assembly.py | 3 ++- cadquery/occ_impl/importers/assembly.py | 12 ++++++------ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/cadquery/assembly.py b/cadquery/assembly.py index 29ec18698..8dda14f37 100644 --- a/cadquery/assembly.py +++ b/cadquery/assembly.py @@ -574,7 +574,8 @@ def importStep(path: str) -> "Assembly": :return: An Assembly object. """ - assy = cast(Assembly, importStepTopLevel(path)) + assy = Assembly() + importStepTopLevel(assy, path) return assy diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 410ece33c..cbdaa8b70 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -10,13 +10,15 @@ from ..assembly import AssemblyProtocol -def importStep(path: str) -> AssemblyProtocol: +def importStep(assy: AssemblyProtocol, path: str): """ Import a step file into an assembly. - """ - # The assembly that is being built from the step file - assy = cq.Assembly() + :param assy: An Assembly object that will be packed with the contents of the STEP file. + :param path: Path and filename to the STEP file to read. + + :return: None + """ # Document that the step file will be read into doc = TDocStd_Document(TCollection_ExtendedString("XmlOcaf")) @@ -115,5 +117,3 @@ def process_label(label, parent_location=None): else: raise ValueError("Step file does not contain an assembly") - - return assy From 4045c5813708dbbd497f2504c5ece94927b8a689 Mon Sep 17 00:00:00 2001 From: AU Date: Mon, 3 Mar 2025 08:07:57 +0100 Subject: [PATCH 06/44] Use Self and @classmethod --- cadquery/assembly.py | 31 +++++++++++++++---------------- 1 file changed, 15 insertions(+), 16 deletions(-) diff --git a/cadquery/assembly.py b/cadquery/assembly.py index 8dda14f37..dfe8cb9ef 100644 --- a/cadquery/assembly.py +++ b/cadquery/assembly.py @@ -11,7 +11,7 @@ cast, get_args, ) -from typing_extensions import Literal +from typing_extensions import Literal, Self from typish import instance_of from uuid import uuid1 as uuid @@ -79,7 +79,6 @@ def _define_grammar(): _grammar = _define_grammar() - class Assembly(object): """Nested assembly of Workplane and Shape objects defining their relative positions.""" @@ -140,7 +139,7 @@ def __init__( self._solve_result = None - def _copy(self) -> "Assembly": + def _copy(self) -> Self: """ Make a deep copy of an assembly """ @@ -164,7 +163,7 @@ def add( loc: Optional[Location] = None, name: Optional[str] = None, color: Optional[Color] = None, - ) -> "Assembly": + ) -> Self: """ Add a subassembly to the current assembly. @@ -186,7 +185,7 @@ def add( name: Optional[str] = None, color: Optional[Color] = None, metadata: Optional[Dict[str, Any]] = None, - ) -> "Assembly": + ) -> Self: """ Add a subassembly to the current assembly with explicit location and name. @@ -299,11 +298,11 @@ def _subloc(self, name: str) -> Tuple[Location, str]: @overload def constrain( self, q1: str, q2: str, kind: ConstraintKind, param: Any = None - ) -> "Assembly": + ) -> Self: ... @overload - def constrain(self, q1: str, kind: ConstraintKind, param: Any = None) -> "Assembly": + def constrain(self, q1: str, kind: ConstraintKind, param: Any = None) -> Self: ... @overload @@ -315,13 +314,13 @@ def constrain( s2: Shape, kind: ConstraintKind, param: Any = None, - ) -> "Assembly": + ) -> Self: ... @overload def constrain( self, id1: str, s1: Shape, kind: ConstraintKind, param: Any = None, - ) -> "Assembly": + ) -> Self: ... def constrain(self, *args, param=None): @@ -366,7 +365,7 @@ def constrain(self, *args, param=None): return self - def solve(self, verbosity: int = 0) -> "Assembly": + def solve(self, verbosity: int = 0) -> Self: """ Solve the constraints. """ @@ -461,7 +460,7 @@ def save( tolerance: float = 0.1, angularTolerance: float = 0.1, **kwargs, - ) -> "Assembly": + ) -> Self: """ Save assembly to a file. @@ -517,7 +516,7 @@ def export( tolerance: float = 0.1, angularTolerance: float = 0.1, **kwargs, - ) -> "Assembly": + ) -> Self: """ Save assembly to a file. @@ -565,8 +564,8 @@ def export( return self - @staticmethod - def importStep(path: str) -> "Assembly": + @classmethod + def importStep(cls, path: str) -> Self: """ Reads an assembly from a STEP file. @@ -574,13 +573,13 @@ def importStep(path: str) -> "Assembly": :return: An Assembly object. """ - assy = Assembly() + assy = cls() importStepTopLevel(assy, path) return assy @classmethod - def load(cls, path: str) -> "Assembly": + def load(cls, path: str) -> Self: raise NotImplementedError From d708acd858480eee6e5b6e79c7c674a965c1791a Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Mon, 3 Mar 2025 07:15:04 -0500 Subject: [PATCH 07/44] Lint fix --- cadquery/assembly.py | 1 + 1 file changed, 1 insertion(+) diff --git a/cadquery/assembly.py b/cadquery/assembly.py index dfe8cb9ef..0b79f6420 100644 --- a/cadquery/assembly.py +++ b/cadquery/assembly.py @@ -79,6 +79,7 @@ def _define_grammar(): _grammar = _define_grammar() + class Assembly(object): """Nested assembly of Workplane and Shape objects defining their relative positions.""" From f88ad3597dcac37d5e725a3646a5ec97fbfa553b Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Tue, 18 Mar 2025 15:25:32 -0400 Subject: [PATCH 08/44] Added a test specifically for testing metadata --- tests/test_assembly.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 1d7b7c387..8977c0e38 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -653,6 +653,21 @@ def test_step_import(tmp_path_factory): assert assy.children[1].loc.toTuple()[0] == (15, 15, 15) +def test_assembly_meta_step_import(tmp_path_factory): + """ + Test import of an assembly with metadata from a STEP file. + """ + + # Use a temporary directory + tmpdir = tmp_path_factory.mktemp("out") + metadata_path = os.path.join(tmpdir, "metadata.step") + + assy = cq.Assembly.importStep(path=metadata_path) + + # Make sure we got the correct number of children + assert len(assy.children) == 6 + + @pytest.mark.parametrize( "assy_fixture, expected", [ From e51a7cbcc2ba5ca2d05f2a1307ae42be1bbbb7a4 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Fri, 18 Apr 2025 09:50:19 -0400 Subject: [PATCH 09/44] Seeing if coverage and tests will pass --- tests/test_assembly.py | 62 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 0a46d414d..ea36cfc76 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -366,6 +366,41 @@ def chassis0_assy(): return chassis +def subshape_assy(): + """ + Builds an assembly with the needed subshapes to test the export and import of STEP files. + """ + + # Create a simple assembly + assy = cq.Assembly(name="top-level") + cube_1 = cq.Workplane().box(10.0, 10.0, 10.0) + assy.add(cube_1, name="cube_1", color=cq.Color("green")) + + # Add subshape name, color and layer + assy.addSubshape( + cube_1.faces(">Z").val(), + name="cube_1_top_face", + color=cq.Color("red"), + layer="cube_1_top_face", + ) + + # Add a cylinder to the assembly + cyl_1 = cq.Workplane().cylinder(10.0, 2.5) + assy.add( + cyl_1, name="cyl_1", color=cq.Color("blue"), loc=cq.Location((0.0, 0.0, -10.0)) + ) + + # Add a subshape face for the cylinder + assy.addSubshape( + cyl_1.faces(" TDocStd_Document: """Read STEP file, return XCAF document""" @@ -785,6 +820,33 @@ def test_meta_step_export_edge_cases(tmp_path_factory): assert success +def test_assembly_step_import(tmp_path_factory): + """ + Test if the STEP import works correctly for an assembly with subshape data attached. + """ + assy = subshape_assy() + + # Use a temporary directory + tmpdir = tmp_path_factory.mktemp("out") + assy_step_path = os.path.join(tmpdir, "assembly_with_subshapes.step") + + success = exportStepMeta(assy, assy_step_path) + assert success + + # Import the STEP file back in + imported_assy = cq.Assembly.importStep(assy_step_path) + + # Check that the assembly was imported successfully + assert imported_assy is not None + + # Check for appropriate part names and colors + # assert imported_assy.children[0].name == "cube_1" + assert imported_assy.children[0].color.toTuple() == (0.0, 1.0, 0.0, 1.0) + # assert imported_assy.children[1].name == "cyl_2" + assert imported_assy.children[1].color.toTuple() == (0.0, 0.0, 1.0, 1.0) + # assert imported_assy.name == "top-level" + + @pytest.mark.parametrize( "assy_fixture, expected", [ From 36d70356f4373a968c2175d88fa8899b052bb4a8 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Mon, 21 Apr 2025 15:18:08 -0400 Subject: [PATCH 10/44] Added name loading from STEP file --- cadquery/occ_impl/importers/assembly.py | 17 +++++++++++++++-- tests/test_assembly.py | 6 +++--- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index cbdaa8b70..757d78b27 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -5,6 +5,7 @@ from OCP.STEPCAFControl import STEPCAFControl_Reader from OCP.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ColorGen, XCAFDoc_ColorSurf from OCP.TDF import TDF_Label, TDF_LabelSequence +from OCP.TDataStd import TDataStd_Name import cadquery as cq from ..assembly import AssemblyProtocol @@ -85,6 +86,12 @@ def process_label(label, parent_location=None): if shape_tool.IsSimpleShape_s(label): shape = shape_tool.GetShape_s(label) + # Load the name of the part in the assembly, if it is present + name = None + name_attr = TDataStd_Name() + if label.FindAttribute(TDataStd_Name.GetID_s(), name_attr): + name = str(name_attr.Get().ToExtString()) + # Process the color for the shape, which could be of different types color = Quantity_Color() if color_tool.GetColor_s(label, XCAFDoc_ColorSurf, color): @@ -102,9 +109,11 @@ def process_label(label, parent_location=None): # Handle the location if it was passed down form a parent component if parent_location is not None: - assy.add(cq.Shape.cast(shape), color=cq_color, loc=parent_location) + assy.add( + cq.Shape.cast(shape), name=name, color=cq_color, loc=parent_location + ) else: - assy.add(cq.Shape.cast(shape), color=cq_color) + assy.add(cq.Shape.cast(shape), name=name, color=cq_color) # Grab the labels, which should hold the assembly parent labels = TDF_LabelSequence() @@ -115,5 +124,9 @@ def process_label(label, parent_location=None): # Start the recursive processing of the assembly process_label(labels.Value(1)) + # Load the top-level name of the assembly, if it is present + name_attr = TDataStd_Name() + if labels.Value(1).FindAttribute(TDataStd_Name.GetID_s(), name_attr): + assy.name = str(name_attr.Get().ToExtString()) else: raise ValueError("Step file does not contain an assembly") diff --git a/tests/test_assembly.py b/tests/test_assembly.py index ea36cfc76..23473c2d6 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -840,11 +840,11 @@ def test_assembly_step_import(tmp_path_factory): assert imported_assy is not None # Check for appropriate part names and colors - # assert imported_assy.children[0].name == "cube_1" + assert imported_assy.children[0].name == "cube_1" assert imported_assy.children[0].color.toTuple() == (0.0, 1.0, 0.0, 1.0) - # assert imported_assy.children[1].name == "cyl_2" + assert imported_assy.children[1].name == "cyl_1" assert imported_assy.children[1].color.toTuple() == (0.0, 0.0, 1.0, 1.0) - # assert imported_assy.name == "top-level" + assert imported_assy.name == "top-level" @pytest.mark.parametrize( From f29056e62a009f24c4b85218ba36e3f8ddbc1b36 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Tue, 22 Apr 2025 12:59:26 -0400 Subject: [PATCH 11/44] Made name a public property of Assembly --- cadquery/occ_impl/assembly.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/cadquery/occ_impl/assembly.py b/cadquery/occ_impl/assembly.py index f2224634b..2c149efe3 100644 --- a/cadquery/occ_impl/assembly.py +++ b/cadquery/occ_impl/assembly.py @@ -140,6 +140,10 @@ def loc(self, value: Location) -> None: def name(self) -> str: ... + @name.setter + def name(self, value: str) -> None: + ... + @property def parent(self) -> Optional["AssemblyProtocol"]: ... From 687bc2b230115bf75246c4344d9c429f95fd9b49 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Tue, 29 Apr 2025 12:40:03 -0400 Subject: [PATCH 12/44] Trying to increase test coverage a bit --- cadquery/occ_impl/importers/assembly.py | 3 +-- tests/test_assembly.py | 9 +++++++++ 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 757d78b27..4bd7e0c2d 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -63,6 +63,7 @@ def process_label(label, parent_location=None): sub_label = comp_labels.Value(i + 1) # The component level holds the location for its shapes + location = parent_location loc = shape_tool.GetLocation_s(sub_label) if loc: location = cq.Location(loc) @@ -77,8 +78,6 @@ def process_label(label, parent_location=None): 0, ): location = parent_location - else: - location = parent_location process_label(sub_label, location) diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 23473c2d6..7bec8c540 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -846,6 +846,15 @@ def test_assembly_step_import(tmp_path_factory): assert imported_assy.children[1].color.toTuple() == (0.0, 0.0, 1.0, 1.0) assert imported_assy.name == "top-level" + # Test a STEP file that does not contain an assembly + wp_step_path = os.path.join(tmpdir, "plain_workplane.step") + res = cq.Workplane().box(10, 10, 10) + res.export(wp_step_path) + + # Import the STEP file back in + with pytest.raises(ValueError): + imported_assy = cq.Assembly.importStep(wp_step_path) + @pytest.mark.parametrize( "assy_fixture, expected", From 6b360c2d56494182e00f00be0855ed3e61f486c6 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Tue, 13 May 2025 11:44:47 -0400 Subject: [PATCH 13/44] Syncing up some experiments --- cadquery/occ_impl/importers/assembly.py | 43 +++++++++++++++++++++++-- tests/test_assembly.py | 32 ++++++++++++++++++ 2 files changed, 73 insertions(+), 2 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 4bd7e0c2d..5ea88019c 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -3,12 +3,13 @@ from OCP.TDocStd import TDocStd_Document from OCP.IFSelect import IFSelect_RetDone from OCP.STEPCAFControl import STEPCAFControl_Reader -from OCP.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ColorGen, XCAFDoc_ColorSurf -from OCP.TDF import TDF_Label, TDF_LabelSequence +from OCP.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ColorGen, XCAFDoc_ColorSurf, XCAFDoc_GraphNode +from OCP.TDF import TDF_Label, TDF_LabelSequence, TDF_AttributeIterator from OCP.TDataStd import TDataStd_Name import cadquery as cq from ..assembly import AssemblyProtocol +from tkinter.constants import CURRENT def importStep(assy: AssemblyProtocol, path: str): @@ -114,6 +115,44 @@ def process_label(label, parent_location=None): else: assy.add(cq.Shape.cast(shape), name=name, color=cq_color) + # if label.NbChildren() > 0: + # child_label = label.FindChild(1) + + # # Create an attribute iterator + # attr_iterator = TDF_AttributeIterator(child_label) + + # # Iterate through all attributes + # while attr_iterator.More(): + # current_attr = attr_iterator.Value() + # # Get the ID of the attribute + # attr_id = current_attr.ID() + # print(f"Found attribute with ID: {attr_id}") + # print(f"Attribute type: {current_attr.DynamicType().Name()}") + # if current_attr.DynamicType().Name() == "XCAFDoc_GraphNode": + # graph_node = XCAFDoc_GraphNode() + # if child_label.FindAttribute(XCAFDoc_GraphNode.GetID_s(), graph_node): + # # Follow the graph node to its referenced labels + # for i in range(1, graph_node.NbChildren() + 1): + # child_graph_node = graph_node.GetChild(i) + # if not child_graph_node.IsNull(): + # referenced_label = child_graph_node.Label() + # if not referenced_label.IsNull(): + # # Check for name on the referenced label + # name_attr = TDataStd_Name() + # if referenced_label.FindAttribute(TDataStd_Name.GetID_s(), name_attr): + # name = name_attr.Get().ToExtString() + # print(f"Found name via GraphNode reference: {name}") + # if current_attr.DynamicType().Name() == "TNaming_NamedShape": + # shape = current_attr.Get() + # if not shape.IsNull(): + # name = shape_tool.GetName_s(shape) + # if name: + # print(f"Shape name: {name}") + # if child_label.FindAttribute(TDataStd_Name.GetID_s(), name_attr): + # print(name_attr) + # Move to next attribute + # attr_iterator.Next() + # Grab the labels, which should hold the assembly parent labels = TDF_LabelSequence() shape_tool.GetFreeShapes(labels) diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 7bec8c540..41ee8804e 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -856,6 +856,38 @@ def test_assembly_step_import(tmp_path_factory): imported_assy = cq.Assembly.importStep(wp_step_path) +def test_assembly_subshape_step_import(tmpdir): + """ + Test if a STEP file containing subshape information can be imported correctly. + """ + + assy_step_path = os.path.join(tmpdir, "subshape_assy.step") + + # Create a basic assembly + cube_1 = cq.Workplane().box(10, 10, 10) + assy = cq.Assembly(name="top_level") + assy.add(cube_1, name="cube_1") + + # Add subshape name, color and layer + assy.addSubshape( + cube_1.faces(">Z").val(), + name="cube_1_top_face", + color=cq.Color("red"), + layer="cube_1_top_face" + ) + + # Export the assembly + success = exportStepMeta(assy, assy_step_path) + assert success + + # Import the STEP file back in + imported_assy = cq.Assembly.importStep(assy_step_path) + assert imported_assy.name == "top_level" + assert len(imported_assy._subshape_names) == 1 + # assert imported_assy.subshapes["cube_1_top_face"].name == "cube_1_top_face" + # assert imported_assy.subshapes["cube_1_top_face"].color == cq.Color("red") + # assert imported_assy.subshapes["cube_1_top_face"].layer == "cube_1_top_face" + @pytest.mark.parametrize( "assy_fixture, expected", [ From 5e193323e3be34abfdde501657900f7a1b160e55 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Thu, 22 May 2025 15:43:22 -0400 Subject: [PATCH 14/44] Got color and layer search working, still need to get name search working through indirect lookup --- cadquery/occ_impl/importers/assembly.py | 83 +++++++++++++------------ tests/test_assembly.py | 5 +- 2 files changed, 47 insertions(+), 41 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 5ea88019c..d56671d3a 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -3,9 +3,14 @@ from OCP.TDocStd import TDocStd_Document from OCP.IFSelect import IFSelect_RetDone from OCP.STEPCAFControl import STEPCAFControl_Reader -from OCP.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ColorGen, XCAFDoc_ColorSurf, XCAFDoc_GraphNode +from OCP.XCAFDoc import ( + XCAFDoc_DocumentTool, + XCAFDoc_ColorGen, + XCAFDoc_ColorSurf, + XCAFDoc_GraphNode, +) from OCP.TDF import TDF_Label, TDF_LabelSequence, TDF_AttributeIterator -from OCP.TDataStd import TDataStd_Name +from OCP.TDataStd import TDataStd_Name, TDataStd_TreeNode import cadquery as cq from ..assembly import AssemblyProtocol @@ -43,6 +48,7 @@ def importStep(assy: AssemblyProtocol, path: str): # Shape and color tools for extracting XCAF data shape_tool = XCAFDoc_DocumentTool.ShapeTool_s(doc.Main()) color_tool = XCAFDoc_DocumentTool.ColorTool_s(doc.Main()) + layer_tool = XCAFDoc_DocumentTool.LayerTool_s(doc.Main()) def process_label(label, parent_location=None): """ @@ -115,43 +121,42 @@ def process_label(label, parent_location=None): else: assy.add(cq.Shape.cast(shape), name=name, color=cq_color) - # if label.NbChildren() > 0: - # child_label = label.FindChild(1) - - # # Create an attribute iterator - # attr_iterator = TDF_AttributeIterator(child_label) - - # # Iterate through all attributes - # while attr_iterator.More(): - # current_attr = attr_iterator.Value() - # # Get the ID of the attribute - # attr_id = current_attr.ID() - # print(f"Found attribute with ID: {attr_id}") - # print(f"Attribute type: {current_attr.DynamicType().Name()}") - # if current_attr.DynamicType().Name() == "XCAFDoc_GraphNode": - # graph_node = XCAFDoc_GraphNode() - # if child_label.FindAttribute(XCAFDoc_GraphNode.GetID_s(), graph_node): - # # Follow the graph node to its referenced labels - # for i in range(1, graph_node.NbChildren() + 1): - # child_graph_node = graph_node.GetChild(i) - # if not child_graph_node.IsNull(): - # referenced_label = child_graph_node.Label() - # if not referenced_label.IsNull(): - # # Check for name on the referenced label - # name_attr = TDataStd_Name() - # if referenced_label.FindAttribute(TDataStd_Name.GetID_s(), name_attr): - # name = name_attr.Get().ToExtString() - # print(f"Found name via GraphNode reference: {name}") - # if current_attr.DynamicType().Name() == "TNaming_NamedShape": - # shape = current_attr.Get() - # if not shape.IsNull(): - # name = shape_tool.GetName_s(shape) - # if name: - # print(f"Shape name: {name}") - # if child_label.FindAttribute(TDataStd_Name.GetID_s(), name_attr): - # print(name_attr) - # Move to next attribute - # attr_iterator.Next() + if label.NbChildren() > 0: + child_label = label.FindChild(1) + attr_iterator = TDF_AttributeIterator(child_label) + while attr_iterator.More(): + current_attr = attr_iterator.Value() + + # Get the type name of the attribute so that we can decide how to handle it + if current_attr.DynamicType().Name() == "TNaming_NamedShape": + # Save the shape so that we can add it to the subshape data + cur_shape = current_attr.Get() + cur_shape_type = cur_shape.ShapeType() + + # Find the layer name, if there is one set for this shape + layers = TDF_LabelSequence() + layer_tool.GetLayers(child_label, layers) + for i in range(1, layers.Length() + 1): + lbl = layers.Value(i) + name_attr = TDataStd_Name() + lbl.FindAttribute(TDataStd_Name.GetID_s(), name_attr) + + # Extract the layer name for the shape here + layer_name = name_attr.Get().ToExtString() + + # Find the subshape color, if there is one set for this shape + color = Quantity_ColorRGBA() + if color_tool.GetColor(cur_shape, XCAFDoc_ColorSurf, color): + rgb = color.GetRGB() + cq_color = cq.Color( + rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha() + ) + elif current_attr.DynamicType().Name() == "TDataStd_TreeNode": + print("TreeNode") + elif current_attr.DynamicType().Name() == "XCAFDoc_GraphNode": + print("GraphNode") + + attr_iterator.Next() # Grab the labels, which should hold the assembly parent labels = TDF_LabelSequence() diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 41ee8804e..b7c502a58 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -872,8 +872,8 @@ def test_assembly_subshape_step_import(tmpdir): assy.addSubshape( cube_1.faces(">Z").val(), name="cube_1_top_face", - color=cq.Color("red"), - layer="cube_1_top_face" + color=cq.Color("blue"), + layer="cube_1_top_face", ) # Export the assembly @@ -888,6 +888,7 @@ def test_assembly_subshape_step_import(tmpdir): # assert imported_assy.subshapes["cube_1_top_face"].color == cq.Color("red") # assert imported_assy.subshapes["cube_1_top_face"].layer == "cube_1_top_face" + @pytest.mark.parametrize( "assy_fixture, expected", [ From 5f4c4c258d1c33768862a275d31789d7192ea53c Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Fri, 23 May 2025 08:56:58 -0400 Subject: [PATCH 15/44] Got tests working for layer and color info import --- cadquery/occ_impl/importers/assembly.py | 11 ++++++++++- tests/test_assembly.py | 18 +++++++++++++----- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index d56671d3a..6a3d3ad6b 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -131,7 +131,9 @@ def process_label(label, parent_location=None): if current_attr.DynamicType().Name() == "TNaming_NamedShape": # Save the shape so that we can add it to the subshape data cur_shape = current_attr.Get() - cur_shape_type = cur_shape.ShapeType() + + # The shape type can be obtained with the following + # cur_shape_type = cur_shape.ShapeType() # Find the layer name, if there is one set for this shape layers = TDF_LabelSequence() @@ -144,13 +146,20 @@ def process_label(label, parent_location=None): # Extract the layer name for the shape here layer_name = name_attr.Get().ToExtString() + # Add the layer as a subshape entry on the assembly + assy.addSubshape(cur_shape, layer=layer_name) + # Find the subshape color, if there is one set for this shape color = Quantity_ColorRGBA() + # Extract the color, if present on the shape if color_tool.GetColor(cur_shape, XCAFDoc_ColorSurf, color): rgb = color.GetRGB() cq_color = cq.Color( rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha() ) + + # Save the color info via the assembly subshape mechanism + assy.addSubshape(cur_shape, color=cq_color) elif current_attr.DynamicType().Name() == "TDataStd_TreeNode": print("TreeNode") elif current_attr.DynamicType().Name() == "XCAFDoc_GraphNode": diff --git a/tests/test_assembly.py b/tests/test_assembly.py index b7c502a58..ac330e446 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -872,7 +872,7 @@ def test_assembly_subshape_step_import(tmpdir): assy.addSubshape( cube_1.faces(">Z").val(), name="cube_1_top_face", - color=cq.Color("blue"), + color=cq.Color("red"), layer="cube_1_top_face", ) @@ -883,10 +883,18 @@ def test_assembly_subshape_step_import(tmpdir): # Import the STEP file back in imported_assy = cq.Assembly.importStep(assy_step_path) assert imported_assy.name == "top_level" - assert len(imported_assy._subshape_names) == 1 - # assert imported_assy.subshapes["cube_1_top_face"].name == "cube_1_top_face" - # assert imported_assy.subshapes["cube_1_top_face"].color == cq.Color("red") - # assert imported_assy.subshapes["cube_1_top_face"].layer == "cube_1_top_face" + + # Check the advanced face name + # assert len(imported_assy._subshape_names) == 1 + # assert list(imported_assy._subshape_names.values())[0] == "cube_1_top_face" + + # Check the color + color = list(imported_assy._subshape_colors.values())[0].toTuple() + assert color == cq.Color("red").toTuple() + + # Check the layer info + layer_name = list(imported_assy._subshape_layers.values())[0] + assert layer_name == "cube_1_top_face" @pytest.mark.parametrize( From e5bccb29ad2f18a98404bd00c9c894491de6729e Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Tue, 27 May 2025 11:22:57 -0400 Subject: [PATCH 16/44] Got shape name loading to work --- cadquery/occ_impl/importers/assembly.py | 133 ++++++++++++++++-------- tests/test_assembly.py | 7 +- 2 files changed, 93 insertions(+), 47 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 6a3d3ad6b..97af54c33 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -9,8 +9,8 @@ XCAFDoc_ColorSurf, XCAFDoc_GraphNode, ) -from OCP.TDF import TDF_Label, TDF_LabelSequence, TDF_AttributeIterator -from OCP.TDataStd import TDataStd_Name, TDataStd_TreeNode +from OCP.TDF import TDF_Label, TDF_LabelSequence, TDF_AttributeIterator, TDF_DataSet +from OCP.TDataStd import TDataStd_Name import cadquery as cq from ..assembly import AssemblyProtocol @@ -92,6 +92,9 @@ def process_label(label, parent_location=None): if shape_tool.IsSimpleShape_s(label): shape = shape_tool.GetShape_s(label) + # Tracks the RGB color value and whether or not it was found + cq_color = None + # Load the name of the part in the assembly, if it is present name = None name_attr = TDataStd_Name() @@ -122,50 +125,92 @@ def process_label(label, parent_location=None): assy.add(cq.Shape.cast(shape), name=name, color=cq_color) if label.NbChildren() > 0: - child_label = label.FindChild(1) - attr_iterator = TDF_AttributeIterator(child_label) - while attr_iterator.More(): - current_attr = attr_iterator.Value() - - # Get the type name of the attribute so that we can decide how to handle it - if current_attr.DynamicType().Name() == "TNaming_NamedShape": - # Save the shape so that we can add it to the subshape data - cur_shape = current_attr.Get() - - # The shape type can be obtained with the following - # cur_shape_type = cur_shape.ShapeType() - - # Find the layer name, if there is one set for this shape - layers = TDF_LabelSequence() - layer_tool.GetLayers(child_label, layers) - for i in range(1, layers.Length() + 1): - lbl = layers.Value(i) - name_attr = TDataStd_Name() - lbl.FindAttribute(TDataStd_Name.GetID_s(), name_attr) - - # Extract the layer name for the shape here - layer_name = name_attr.Get().ToExtString() - - # Add the layer as a subshape entry on the assembly - assy.addSubshape(cur_shape, layer=layer_name) - - # Find the subshape color, if there is one set for this shape - color = Quantity_ColorRGBA() - # Extract the color, if present on the shape - if color_tool.GetColor(cur_shape, XCAFDoc_ColorSurf, color): - rgb = color.GetRGB() - cq_color = cq.Color( - rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha() + for j in range(label.NbChildren()): + child_label = label.FindChild(j + 1) + attr_iterator = TDF_AttributeIterator(child_label) + while attr_iterator.More(): + current_attr = attr_iterator.Value() + + # Get the type name of the attribute so that we can decide how to handle it + if current_attr.DynamicType().Name() == "TNaming_NamedShape": + # Save the shape so that we can add it to the subshape data + cur_shape = current_attr.Get() + + # Find the layer name, if there is one set for this shape + layers = TDF_LabelSequence() + layer_tool.GetLayers(child_label, layers) + for i in range(1, layers.Length() + 1): + lbl = layers.Value(i) + name_attr = TDataStd_Name() + lbl.FindAttribute(TDataStd_Name.GetID_s(), name_attr) + + # Extract the layer name for the shape here + layer_name = name_attr.Get().ToExtString() + + # Add the layer as a subshape entry on the assembly + assy.addSubshape(cur_shape, layer=layer_name) + + # Find the subshape color, if there is one set for this shape + color = Quantity_ColorRGBA() + # Extract the color, if present on the shape + if color_tool.GetColor(cur_shape, XCAFDoc_ColorSurf, color): + rgb = color.GetRGB() + cq_color = cq.Color( + rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha() + ) + + # Save the color info via the assembly subshape mechanism + assy.addSubshape(cur_shape, color=cq_color) + elif current_attr.DynamicType().Name() == "TDataStd_TreeNode": + # Holds the color name, if found, and tells us whether or not it was found + color_name = None + + # Get the attributes of the father node + father_attr = current_attr.Father() + + # Iterate theough the attributes to see if there is a color name + lbl = father_attr.Label() + it = TDF_AttributeIterator(lbl) + while it.More(): + new_attr = it.Value() + if new_attr.DynamicType().Name() == "TDataStd_Name": + # Retrieve the name + name_string = new_attr.Get().ToExtString() + + # Make sure that we have a color name + if "#" in name_string: + color_name = name_string.split(" ")[0] + + it.Next() + + # If we found a color name, save it on the subshape + # Perfer the RGB value because when importing, OCCT will try to round + # RGB values to fit color names. + if cq_color is not None: + assy.addSubshape(cur_shape, color=cq_color) + elif color_name is not None: + assy.addSubshape(cur_shape, color=cq.Color(color_name)) + elif current_attr.DynamicType().Name() == "XCAFDoc_GraphNode": + # Step up one level to try to get the name from the parent + lbl = current_attr.GetFather(1).Label() + + # Step through and search for the name attribute + it = TDF_AttributeIterator(lbl) + while it.More(): + new_attr = it.Value() + if new_attr.DynamicType().Name() == "TDataStd_Name": + # Save this as the name of the subshape + assy.addSubshape( + cur_shape, name=new_attr.Get().ToExtString(), + ) + it.Next() + else: + print( + "Unknown attribute type:", + current_attr.DynamicType().Name(), ) - # Save the color info via the assembly subshape mechanism - assy.addSubshape(cur_shape, color=cq_color) - elif current_attr.DynamicType().Name() == "TDataStd_TreeNode": - print("TreeNode") - elif current_attr.DynamicType().Name() == "XCAFDoc_GraphNode": - print("GraphNode") - - attr_iterator.Next() + attr_iterator.Next() # Grab the labels, which should hold the assembly parent labels = TDF_LabelSequence() diff --git a/tests/test_assembly.py b/tests/test_assembly.py index ac330e446..3f5baad12 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -856,11 +856,12 @@ def test_assembly_step_import(tmp_path_factory): imported_assy = cq.Assembly.importStep(wp_step_path) -def test_assembly_subshape_step_import(tmpdir): +def test_assembly_subshape_step_import(tmp_path_factory): """ Test if a STEP file containing subshape information can be imported correctly. """ + tmpdir = tmp_path_factory.mktemp("out") assy_step_path = os.path.join(tmpdir, "subshape_assy.step") # Create a basic assembly @@ -885,8 +886,8 @@ def test_assembly_subshape_step_import(tmpdir): assert imported_assy.name == "top_level" # Check the advanced face name - # assert len(imported_assy._subshape_names) == 1 - # assert list(imported_assy._subshape_names.values())[0] == "cube_1_top_face" + assert len(imported_assy._subshape_names) == 1 + assert list(imported_assy._subshape_names.values())[0] == "cube_1_top_face" # Check the color color = list(imported_assy._subshape_colors.values())[0].toTuple() From d5bc317318e121145b1a804300be181c991f41fa Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Tue, 27 May 2025 15:38:34 -0400 Subject: [PATCH 17/44] Trying to get approximate tuple comparison working --- tests/test_assembly.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 3f5baad12..e916c238e 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -839,11 +839,14 @@ def test_assembly_step_import(tmp_path_factory): # Check that the assembly was imported successfully assert imported_assy is not None - # Check for appropriate part names and colors + # Check for appropriate part name assert imported_assy.children[0].name == "cube_1" - assert imported_assy.children[0].color.toTuple() == (0.0, 1.0, 0.0, 1.0) + # Check for approximate color match + assert pytest.approx(imported_assy.children[0].color.toTuple(), rel=0.01) == (0.0, 1.0, 0.0, 1.0) + # Check for appropriate part name assert imported_assy.children[1].name == "cyl_1" - assert imported_assy.children[1].color.toTuple() == (0.0, 0.0, 1.0, 1.0) + # Check for approximate color match + assert pytest.approx(imported_assy.children[1].color.toTuple(), rel=0.01) == (0.0, 0.0, 1.0, 1.0) assert imported_assy.name == "top-level" # Test a STEP file that does not contain an assembly From 9b48507e5341ecb4d9cc6edb30a8e55c6dbf7acf Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Tue, 27 May 2025 16:01:19 -0400 Subject: [PATCH 18/44] Black fix --- tests/test_assembly.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tests/test_assembly.py b/tests/test_assembly.py index e916c238e..73689c3f6 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -842,11 +842,21 @@ def test_assembly_step_import(tmp_path_factory): # Check for appropriate part name assert imported_assy.children[0].name == "cube_1" # Check for approximate color match - assert pytest.approx(imported_assy.children[0].color.toTuple(), rel=0.01) == (0.0, 1.0, 0.0, 1.0) + assert pytest.approx(imported_assy.children[0].color.toTuple(), rel=0.01) == ( + 0.0, + 1.0, + 0.0, + 1.0, + ) # Check for appropriate part name assert imported_assy.children[1].name == "cyl_1" # Check for approximate color match - assert pytest.approx(imported_assy.children[1].color.toTuple(), rel=0.01) == (0.0, 0.0, 1.0, 1.0) + assert pytest.approx(imported_assy.children[1].color.toTuple(), rel=0.01) == ( + 0.0, + 0.0, + 1.0, + 1.0, + ) assert imported_assy.name == "top-level" # Test a STEP file that does not contain an assembly From 65f38c87f9d6d6a5741976f5524d18c9a13d8b79 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Wed, 28 May 2025 08:34:20 -0400 Subject: [PATCH 19/44] Added a test for a bad filename, and added a custom color --- tests/test_assembly.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 73689c3f6..4cfe48120 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -880,7 +880,7 @@ def test_assembly_subshape_step_import(tmp_path_factory): # Create a basic assembly cube_1 = cq.Workplane().box(10, 10, 10) assy = cq.Assembly(name="top_level") - assy.add(cube_1, name="cube_1") + assy.add(cube_1, name="cube_1", color=cq.Color(0.76512, 0.23491, 0.91301)) # Add subshape name, color and layer assy.addSubshape( @@ -911,6 +911,20 @@ def test_assembly_subshape_step_import(tmp_path_factory): assert layer_name == "cube_1_top_face" +def test_bad_step_file_import(tmp_path_factory): + """ + Test if a bad STEP file raises an error when importing. + """ + + tmpdir = tmp_path_factory.mktemp("out") + bad_step_path = os.path.join(tmpdir, "bad_step.step") + + # Check that an error is raised when trying to import a non-existent STEP file + with pytest.raises(ValueError): + # Export the assembly + imported_assy = cq.Assembly.importStep(bad_step_path) + + @pytest.mark.parametrize( "assy_fixture, expected", [ From 2810bc22b325c2890d363f82c9f3f2e06fb7c1c7 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Wed, 28 May 2025 08:56:13 -0400 Subject: [PATCH 20/44] Increase test coverage a bit and improve color name check --- cadquery/occ_impl/importers/assembly.py | 7 +------ tests/test_assembly.py | 8 ++++---- 2 files changed, 5 insertions(+), 10 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 97af54c33..3c29baf40 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -184,7 +184,7 @@ def process_label(label, parent_location=None): it.Next() # If we found a color name, save it on the subshape - # Perfer the RGB value because when importing, OCCT will try to round + # Perfer the RGB value because when importing, OCCT tries to round # RGB values to fit color names. if cq_color is not None: assy.addSubshape(cur_shape, color=cq_color) @@ -204,11 +204,6 @@ def process_label(label, parent_location=None): cur_shape, name=new_attr.Get().ToExtString(), ) it.Next() - else: - print( - "Unknown attribute type:", - current_attr.DynamicType().Name(), - ) attr_iterator.Next() diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 5bcea5927..3357aa267 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -34,7 +34,7 @@ from OCP.STEPCAFControl import STEPCAFControl_Reader from OCP.IFSelect import IFSelect_RetDone from OCP.TDF import TDF_ChildIterator -from OCP.Quantity import Quantity_ColorRGBA, Quantity_TOC_sRGB +from OCP.Quantity import Quantity_ColorRGBA, Quantity_TOC_sRGB, Quantity_NameOfColor from OCP.TopAbs import TopAbs_ShapeEnum @@ -880,7 +880,7 @@ def test_assembly_subshape_step_import(tmp_path_factory): # Create a basic assembly cube_1 = cq.Workplane().box(10, 10, 10) assy = cq.Assembly(name="top_level") - assy.add(cube_1, name="cube_1", color=cq.Color(0.76512, 0.23491, 0.91301)) + assy.add(cube_1, name="cube_1") # Add subshape name, color and layer assy.addSubshape( @@ -903,8 +903,8 @@ def test_assembly_subshape_step_import(tmp_path_factory): assert list(imported_assy._subshape_names.values())[0] == "cube_1_top_face" # Check the color - color = list(imported_assy._subshape_colors.values())[0].toTuple() - assert color == cq.Color("red").toTuple() + color = list(imported_assy._subshape_colors.values())[0] + assert Quantity_NameOfColor.Quantity_NOC_RED == color.wrapped.GetRGB().Name() # Check the layer info layer_name = list(imported_assy._subshape_layers.values())[0] From 3438df71a3c088008db337f14e768d131d04e1fa Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Wed, 28 May 2025 17:25:01 -0400 Subject: [PATCH 21/44] Removing code that should never be hit --- cadquery/occ_impl/importers/assembly.py | 37 +------------------------ 1 file changed, 1 insertion(+), 36 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 3c29baf40..649db1e93 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -103,18 +103,12 @@ def process_label(label, parent_location=None): # Process the color for the shape, which could be of different types color = Quantity_Color() + cq_color = cq.Color(0.0, 0.0, 0.0) if color_tool.GetColor_s(label, XCAFDoc_ColorSurf, color): r = color.Red() g = color.Green() b = color.Blue() cq_color = cq.Color(r, g, b) - elif color_tool.GetColor_s(label, XCAFDoc_ColorGen, color): - r = color.Red() - g = color.Green() - b = color.Blue() - cq_color = cq.Color(r, g, b) - else: - cq_color = cq.Color(0.5, 0.5, 0.5) # Handle the location if it was passed down form a parent component if parent_location is not None: @@ -161,35 +155,6 @@ def process_label(label, parent_location=None): # Save the color info via the assembly subshape mechanism assy.addSubshape(cur_shape, color=cq_color) - elif current_attr.DynamicType().Name() == "TDataStd_TreeNode": - # Holds the color name, if found, and tells us whether or not it was found - color_name = None - - # Get the attributes of the father node - father_attr = current_attr.Father() - - # Iterate theough the attributes to see if there is a color name - lbl = father_attr.Label() - it = TDF_AttributeIterator(lbl) - while it.More(): - new_attr = it.Value() - if new_attr.DynamicType().Name() == "TDataStd_Name": - # Retrieve the name - name_string = new_attr.Get().ToExtString() - - # Make sure that we have a color name - if "#" in name_string: - color_name = name_string.split(" ")[0] - - it.Next() - - # If we found a color name, save it on the subshape - # Perfer the RGB value because when importing, OCCT tries to round - # RGB values to fit color names. - if cq_color is not None: - assy.addSubshape(cur_shape, color=cq_color) - elif color_name is not None: - assy.addSubshape(cur_shape, color=cq.Color(color_name)) elif current_attr.DynamicType().Name() == "XCAFDoc_GraphNode": # Step up one level to try to get the name from the parent lbl = current_attr.GetFather(1).Label() From 086fa8ac7d190961197bd0bb0e1674fdc39a049d Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Thu, 29 May 2025 14:44:07 -0400 Subject: [PATCH 22/44] Still trying to increase test coverage --- cadquery/occ_impl/importers/assembly.py | 111 ++++++++++++------------ 1 file changed, 56 insertions(+), 55 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 649db1e93..0f4596c5f 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -72,6 +72,7 @@ def process_label(label, parent_location=None): # The component level holds the location for its shapes location = parent_location loc = shape_tool.GetLocation_s(sub_label) + location = cq.Location((0.0, 0.0, 0.0)) if loc: location = cq.Location(loc) @@ -98,8 +99,8 @@ def process_label(label, parent_location=None): # Load the name of the part in the assembly, if it is present name = None name_attr = TDataStd_Name() - if label.FindAttribute(TDataStd_Name.GetID_s(), name_attr): - name = str(name_attr.Get().ToExtString()) + label.FindAttribute(TDataStd_Name.GetID_s(), name_attr) + name = str(name_attr.Get().ToExtString()) # Process the color for the shape, which could be of different types color = Quantity_Color() @@ -118,59 +119,59 @@ def process_label(label, parent_location=None): else: assy.add(cq.Shape.cast(shape), name=name, color=cq_color) - if label.NbChildren() > 0: - for j in range(label.NbChildren()): - child_label = label.FindChild(j + 1) - attr_iterator = TDF_AttributeIterator(child_label) - while attr_iterator.More(): - current_attr = attr_iterator.Value() - - # Get the type name of the attribute so that we can decide how to handle it - if current_attr.DynamicType().Name() == "TNaming_NamedShape": - # Save the shape so that we can add it to the subshape data - cur_shape = current_attr.Get() - - # Find the layer name, if there is one set for this shape - layers = TDF_LabelSequence() - layer_tool.GetLayers(child_label, layers) - for i in range(1, layers.Length() + 1): - lbl = layers.Value(i) - name_attr = TDataStd_Name() - lbl.FindAttribute(TDataStd_Name.GetID_s(), name_attr) - - # Extract the layer name for the shape here - layer_name = name_attr.Get().ToExtString() - - # Add the layer as a subshape entry on the assembly - assy.addSubshape(cur_shape, layer=layer_name) - - # Find the subshape color, if there is one set for this shape - color = Quantity_ColorRGBA() - # Extract the color, if present on the shape - if color_tool.GetColor(cur_shape, XCAFDoc_ColorSurf, color): - rgb = color.GetRGB() - cq_color = cq.Color( - rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha() + # Check all the attributes of all the children to find the subshapes and any names + for j in range(label.NbChildren()): + child_label = label.FindChild(j + 1) + attr_iterator = TDF_AttributeIterator(child_label) + while attr_iterator.More(): + current_attr = attr_iterator.Value() + + # Get the type name of the attribute so that we can decide how to handle it + if current_attr.DynamicType().Name() == "TNaming_NamedShape": + # Save the shape so that we can add it to the subshape data + cur_shape = current_attr.Get() + + # Find the layer name, if there is one set for this shape + layers = TDF_LabelSequence() + layer_tool.GetLayers(child_label, layers) + for i in range(1, layers.Length() + 1): + lbl = layers.Value(i) + name_attr = TDataStd_Name() + lbl.FindAttribute(TDataStd_Name.GetID_s(), name_attr) + + # Extract the layer name for the shape here + layer_name = name_attr.Get().ToExtString() + + # Add the layer as a subshape entry on the assembly + assy.addSubshape(cur_shape, layer=layer_name) + + # Find the subshape color, if there is one set for this shape + color = Quantity_ColorRGBA() + # Extract the color, if present on the shape + if color_tool.GetColor(cur_shape, XCAFDoc_ColorSurf, color): + rgb = color.GetRGB() + cq_color = cq.Color( + rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha() + ) + + # Save the color info via the assembly subshape mechanism + assy.addSubshape(cur_shape, color=cq_color) + elif current_attr.DynamicType().Name() == "XCAFDoc_GraphNode": + # Step up one level to try to get the name from the parent + lbl = current_attr.GetFather(1).Label() + + # Step through and search for the name attribute + it = TDF_AttributeIterator(lbl) + while it.More(): + new_attr = it.Value() + if new_attr.DynamicType().Name() == "TDataStd_Name": + # Save this as the name of the subshape + assy.addSubshape( + cur_shape, name=new_attr.Get().ToExtString(), ) + it.Next() - # Save the color info via the assembly subshape mechanism - assy.addSubshape(cur_shape, color=cq_color) - elif current_attr.DynamicType().Name() == "XCAFDoc_GraphNode": - # Step up one level to try to get the name from the parent - lbl = current_attr.GetFather(1).Label() - - # Step through and search for the name attribute - it = TDF_AttributeIterator(lbl) - while it.More(): - new_attr = it.Value() - if new_attr.DynamicType().Name() == "TDataStd_Name": - # Save this as the name of the subshape - assy.addSubshape( - cur_shape, name=new_attr.Get().ToExtString(), - ) - it.Next() - - attr_iterator.Next() + attr_iterator.Next() # Grab the labels, which should hold the assembly parent labels = TDF_LabelSequence() @@ -183,7 +184,7 @@ def process_label(label, parent_location=None): # Load the top-level name of the assembly, if it is present name_attr = TDataStd_Name() - if labels.Value(1).FindAttribute(TDataStd_Name.GetID_s(), name_attr): - assy.name = str(name_attr.Get().ToExtString()) + labels.Value(1).FindAttribute(TDataStd_Name.GetID_s(), name_attr) + assy.name = str(name_attr.Get().ToExtString()) else: raise ValueError("Step file does not contain an assembly") From 636551238818a31a105bf9a8a947acf0da664f5e Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Thu, 29 May 2025 15:43:05 -0400 Subject: [PATCH 23/44] Added a test for a plain assembly --- tests/test_assembly.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 3357aa267..8c7c687d6 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -925,6 +925,28 @@ def test_bad_step_file_import(tmp_path_factory): imported_assy = cq.Assembly.importStep(bad_step_path) +def test_plain_assembly_import(tmp_path_factory): + """ + Test to make sure that importing plain assemblies has not been broken. + """ + + tmpdir = tmp_path_factory.mktemp("out") + plain_step_path = os.path.join(tmpdir, "plain_assembly_step.step") + + # Create a basic assembly + cube_1 = cq.Workplane().box(10, 10, 10) + assy = cq.Assembly(name="top_level") + assy.add(cube_1) + + # Export the assembly + success = exportStepMeta(assy, plain_step_path) + assert success + + # Import the STEP file back in + imported_assy = cq.Assembly.importStep(plain_step_path) + assert imported_assy.name == "top_level" + + @pytest.mark.parametrize( "assy_fixture, expected", [ From 578fd60060f01c3930da97b78200327d6fb71b41 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Tue, 3 Jun 2025 09:16:44 -0400 Subject: [PATCH 24/44] Refactored a bit to support nested assemblies better in the future --- cadquery/occ_impl/importers/assembly.py | 229 ++++++++++++------------ 1 file changed, 112 insertions(+), 117 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 0f4596c5f..78fba0fc5 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -50,6 +50,89 @@ def importStep(assy: AssemblyProtocol, path: str): color_tool = XCAFDoc_DocumentTool.ColorTool_s(doc.Main()) layer_tool = XCAFDoc_DocumentTool.LayerTool_s(doc.Main()) + def _process_simple_shape(label, parent_location=None): + shape = shape_tool.GetShape_s(label) + + # Tracks the RGB color value and whether or not it was found + cq_color = None + + # Load the name of the part in the assembly, if it is present + name = None + name_attr = TDataStd_Name() + label.FindAttribute(TDataStd_Name.GetID_s(), name_attr) + name = str(name_attr.Get().ToExtString()) + + # Process the color for the shape, which could be of different types + color = Quantity_Color() + cq_color = cq.Color(0.0, 0.0, 0.0) + if color_tool.GetColor_s(label, XCAFDoc_ColorSurf, color): + r = color.Red() + g = color.Green() + b = color.Blue() + cq_color = cq.Color(r, g, b) + + # Handle the location if it was passed down form a parent component + if parent_location is not None: + assy.add( + cq.Shape.cast(shape), name=name, color=cq_color, loc=parent_location + ) + else: + assy.add(cq.Shape.cast(shape), name=name, color=cq_color) + + # Check all the attributes of all the children to find the subshapes and any names + for j in range(label.NbChildren()): + child_label = label.FindChild(j + 1) + attr_iterator = TDF_AttributeIterator(child_label) + while attr_iterator.More(): + current_attr = attr_iterator.Value() + + # Get the type name of the attribute so that we can decide how to handle it + if current_attr.DynamicType().Name() == "TNaming_NamedShape": + # Save the shape so that we can add it to the subshape data + cur_shape = current_attr.Get() + + # Find the layer name, if there is one set for this shape + layers = TDF_LabelSequence() + layer_tool.GetLayers(child_label, layers) + for i in range(1, layers.Length() + 1): + lbl = layers.Value(i) + name_attr = TDataStd_Name() + lbl.FindAttribute(TDataStd_Name.GetID_s(), name_attr) + + # Extract the layer name for the shape here + layer_name = name_attr.Get().ToExtString() + + # Add the layer as a subshape entry on the assembly + assy.addSubshape(cur_shape, layer=layer_name) + + # Find the subshape color, if there is one set for this shape + color = Quantity_ColorRGBA() + # Extract the color, if present on the shape + if color_tool.GetColor(cur_shape, XCAFDoc_ColorSurf, color): + rgb = color.GetRGB() + cq_color = cq.Color( + rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha() + ) + + # Save the color info via the assembly subshape mechanism + assy.addSubshape(cur_shape, color=cq_color) + elif current_attr.DynamicType().Name() == "XCAFDoc_GraphNode": + # Step up one level to try to get the name from the parent + lbl = current_attr.GetFather(1).Label() + + # Step through and search for the name attribute + it = TDF_AttributeIterator(lbl) + while it.More(): + new_attr = it.Value() + if new_attr.DynamicType().Name() == "TDataStd_Name": + # Save this as the name of the subshape + assy.addSubshape( + cur_shape, name=new_attr.Get().ToExtString(), + ) + it.Next() + + attr_iterator.Next() + def process_label(label, parent_location=None): """ Recursive function that allows us to process the hierarchy of the assembly as represented @@ -61,130 +144,42 @@ def process_label(label, parent_location=None): ref_label = TDF_Label() shape_tool.GetReferredShape_s(label, ref_label) process_label(ref_label, parent_location) + + # Load the top-level name of the assembly, if it is present + name_attr = TDataStd_Name() + labels.Value(1).FindAttribute(TDataStd_Name.GetID_s(), name_attr) + assy.name = str(name_attr.Get().ToExtString()) + return - # Process components - comp_labels = TDF_LabelSequence() - shape_tool.GetComponents_s(label, comp_labels) - for i in range(comp_labels.Length()): - sub_label = comp_labels.Value(i + 1) - - # The component level holds the location for its shapes - location = parent_location - loc = shape_tool.GetLocation_s(sub_label) - location = cq.Location((0.0, 0.0, 0.0)) - if loc: - location = cq.Location(loc) - - # Make sure that the location object is actually doing something interesting - # This is done because the location may have to go through multiple levels of - # components before the shapes are found. This allows the top-level component - # to specify the location/rotation of the shapes. - if location.toTuple()[0] == (0, 0, 0) and location.toTuple()[1] == ( - 0, - 0, - 0, - ): - location = parent_location - - process_label(sub_label, location) - - # Check to see if we have an endpoint shape - if shape_tool.IsSimpleShape_s(label): - shape = shape_tool.GetShape_s(label) + # See if this is an assembly (or sub-assembly) + if shape_tool.IsAssembly_s(label): + # Recursively process its components (children) + comp_labels = TDF_LabelSequence() + shape_tool.GetComponents_s(label, comp_labels) + for i in range(comp_labels.Length()): + sub_label = comp_labels.Value(i + 1) - # Tracks the RGB color value and whether or not it was found - cq_color = None + # Pass down the location or other context as needed + process_label(sub_label, parent_location) + return - # Load the name of the part in the assembly, if it is present - name = None - name_attr = TDataStd_Name() - label.FindAttribute(TDataStd_Name.GetID_s(), name_attr) - name = str(name_attr.Get().ToExtString()) - - # Process the color for the shape, which could be of different types - color = Quantity_Color() - cq_color = cq.Color(0.0, 0.0, 0.0) - if color_tool.GetColor_s(label, XCAFDoc_ColorSurf, color): - r = color.Red() - g = color.Green() - b = color.Blue() - cq_color = cq.Color(r, g, b) - - # Handle the location if it was passed down form a parent component - if parent_location is not None: - assy.add( - cq.Shape.cast(shape), name=name, color=cq_color, loc=parent_location - ) - else: - assy.add(cq.Shape.cast(shape), name=name, color=cq_color) - - # Check all the attributes of all the children to find the subshapes and any names - for j in range(label.NbChildren()): - child_label = label.FindChild(j + 1) - attr_iterator = TDF_AttributeIterator(child_label) - while attr_iterator.More(): - current_attr = attr_iterator.Value() - - # Get the type name of the attribute so that we can decide how to handle it - if current_attr.DynamicType().Name() == "TNaming_NamedShape": - # Save the shape so that we can add it to the subshape data - cur_shape = current_attr.Get() - - # Find the layer name, if there is one set for this shape - layers = TDF_LabelSequence() - layer_tool.GetLayers(child_label, layers) - for i in range(1, layers.Length() + 1): - lbl = layers.Value(i) - name_attr = TDataStd_Name() - lbl.FindAttribute(TDataStd_Name.GetID_s(), name_attr) - - # Extract the layer name for the shape here - layer_name = name_attr.Get().ToExtString() - - # Add the layer as a subshape entry on the assembly - assy.addSubshape(cur_shape, layer=layer_name) - - # Find the subshape color, if there is one set for this shape - color = Quantity_ColorRGBA() - # Extract the color, if present on the shape - if color_tool.GetColor(cur_shape, XCAFDoc_ColorSurf, color): - rgb = color.GetRGB() - cq_color = cq.Color( - rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha() - ) + # Check to see if we have an endpoint shape and process it + if shape_tool.IsSimpleShape_s(label): + _process_simple_shape(label, parent_location) - # Save the color info via the assembly subshape mechanism - assy.addSubshape(cur_shape, color=cq_color) - elif current_attr.DynamicType().Name() == "XCAFDoc_GraphNode": - # Step up one level to try to get the name from the parent - lbl = current_attr.GetFather(1).Label() - - # Step through and search for the name attribute - it = TDF_AttributeIterator(lbl) - while it.More(): - new_attr = it.Value() - if new_attr.DynamicType().Name() == "TDataStd_Name": - # Save this as the name of the subshape - assy.addSubshape( - cur_shape, name=new_attr.Get().ToExtString(), - ) - it.Next() - - attr_iterator.Next() - - # Grab the labels, which should hold the assembly parent + # Look for the top-level assembly + found_top_level_assembly = False labels = TDF_LabelSequence() shape_tool.GetFreeShapes(labels) + for i in range(labels.Length()): + # Make sure that we have an assembly at the top level + if shape_tool.IsTopLevel(labels.Value(i + 1)): + if shape_tool.IsAssembly_s(labels.Value(i + 1)): + found_top_level_assembly = True - # Make sure that we are working with an assembly - if shape_tool.IsAssembly_s(labels.Value(1)): - # Start the recursive processing of the assembly - process_label(labels.Value(1)) + process_label(labels.Value(i + 1)) - # Load the top-level name of the assembly, if it is present - name_attr = TDataStd_Name() - labels.Value(1).FindAttribute(TDataStd_Name.GetID_s(), name_attr) - assy.name = str(name_attr.Get().ToExtString()) - else: + # If we did not find a top-level assembly, raise an error + if not found_top_level_assembly: raise ValueError("Step file does not contain an assembly") From c90e47655f7e34c4bf11387b43d9c1ee0a743ec6 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Tue, 3 Jun 2025 16:23:00 -0400 Subject: [PATCH 25/44] Fixed location handling for components of the assembly --- cadquery/occ_impl/importers/assembly.py | 9 +++- tests/test_assembly.py | 57 +++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 78fba0fc5..5f15baddd 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -74,7 +74,10 @@ def _process_simple_shape(label, parent_location=None): # Handle the location if it was passed down form a parent component if parent_location is not None: assy.add( - cq.Shape.cast(shape), name=name, color=cq_color, loc=parent_location + cq.Shape.cast(shape), + name=name, + color=cq_color, + loc=cq.Location(parent_location), ) else: assy.add(cq.Shape.cast(shape), name=name, color=cq_color) @@ -159,9 +162,11 @@ def process_label(label, parent_location=None): shape_tool.GetComponents_s(label, comp_labels) for i in range(comp_labels.Length()): sub_label = comp_labels.Value(i + 1) + # Get the location of the sub-label, if it exists + loc = shape_tool.GetLocation_s(sub_label) # Pass down the location or other context as needed - process_label(sub_label, parent_location) + process_label(sub_label, loc) return # Check to see if we have an endpoint shape and process it diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 8c7c687d6..bc295a115 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -857,6 +857,11 @@ def test_assembly_step_import(tmp_path_factory): 1.0, 1.0, ) + + # Make sure the shape locations were applied correctly + assert imported_assy.children[1].loc.toTuple()[0] == (0.0, 0.0, -10.0) + + # Check the top-level assembly name assert imported_assy.name == "top-level" # Test a STEP file that does not contain an assembly @@ -911,6 +916,58 @@ def test_assembly_subshape_step_import(tmp_path_factory): assert layer_name == "cube_1_top_face" +def test_assembly_multi_subshape_step_import(tmp_path_factory): + """ + Test if a STEP file containing subshape information can be imported correctly. + """ + + tmpdir = tmp_path_factory.mktemp("out") + assy_step_path = os.path.join(tmpdir, "multi_subshape_assy.step") + + # Create a basic assembly + cube_1 = cq.Workplane().box(10, 10, 10) + assy = cq.Assembly(name="top_level") + assy.add(cube_1, name="cube_1", color=cq.Color("green")) + cube_2 = cq.Workplane().box(5, 5, 5) + assy.add(cube_2, name="cube_2", color=cq.Color("blue"), loc=cq.Location(10, 10, 10)) + + # Add subshape name, color and layer + assy.addSubshape( + cube_1.faces(">Z").val(), + name="cube_1_top_face", + color=cq.Color("red"), + layer="cube_1_top_face", + ) + assy.addSubshape( + cube_2.faces(">X").val(), + name="cube_2_right_face", + color=cq.Color("red"), + layer="cube_2_right_face", + ) + + # Export the assembly + success = exportStepMeta(assy, assy_step_path) + assert success + + # Import the STEP file back in + imported_assy = cq.Assembly.importStep(assy_step_path) + + # Check that the top-level assembly name is correct + assert imported_assy.name == "top_level" + + # Check the advanced face name + assert len(imported_assy._subshape_names) == 2 + assert list(imported_assy._subshape_names.values())[0] == "cube_1_top_face" + + # Check the color + color = list(imported_assy._subshape_colors.values())[0] + assert Quantity_NameOfColor.Quantity_NOC_RED == color.wrapped.GetRGB().Name() + + # Check the layer info + layer_name = list(imported_assy._subshape_layers.values())[0] + assert layer_name == "cube_1_top_face" + + def test_bad_step_file_import(tmp_path_factory): """ Test if a bad STEP file raises an error when importing. From 83a433f05674ba497adfe8d159082adc5081b95c Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Tue, 10 Jun 2025 11:08:05 -0400 Subject: [PATCH 26/44] Fix the default color to not be black --- cadquery/occ_impl/importers/assembly.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 5f15baddd..380facf16 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -64,7 +64,7 @@ def _process_simple_shape(label, parent_location=None): # Process the color for the shape, which could be of different types color = Quantity_Color() - cq_color = cq.Color(0.0, 0.0, 0.0) + cq_color = cq.Color(0.50, 0.50, 0.50) if color_tool.GetColor_s(label, XCAFDoc_ColorSurf, color): r = color.Red() g = color.Green() From cafad47ef2806354b743b7d967d65db53674d816 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Thu, 12 Jun 2025 15:15:28 -0400 Subject: [PATCH 27/44] Fixed bug with parent location not being applied when needed --- cadquery/occ_impl/importers/assembly.py | 4 ++ tests/test_assembly.py | 52 ++++++++++++++++++++++--- 2 files changed, 50 insertions(+), 6 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 380facf16..f03384b78 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -160,12 +160,16 @@ def process_label(label, parent_location=None): # Recursively process its components (children) comp_labels = TDF_LabelSequence() shape_tool.GetComponents_s(label, comp_labels) + for i in range(comp_labels.Length()): sub_label = comp_labels.Value(i + 1) # Get the location of the sub-label, if it exists loc = shape_tool.GetLocation_s(sub_label) # Pass down the location or other context as needed + # Add the parent location if it exists + if parent_location is not None: + loc = parent_location * loc process_label(sub_label, loc) return diff --git a/tests/test_assembly.py b/tests/test_assembly.py index bc295a115..853d5b6c0 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -987,22 +987,62 @@ def test_plain_assembly_import(tmp_path_factory): Test to make sure that importing plain assemblies has not been broken. """ + from cadquery.func import box, rect + tmpdir = tmp_path_factory.mktemp("out") plain_step_path = os.path.join(tmpdir, "plain_assembly_step.step") - # Create a basic assembly + # Simple cubes cube_1 = cq.Workplane().box(10, 10, 10) + cube_2 = cq.Workplane().box(5, 5, 5) + cube_3 = cq.Workplane().box(5, 5, 5) + cube_4 = cq.Workplane().box(5, 5, 5) + assy = cq.Assembly(name="top_level") - assy.add(cube_1) + assy.add(cube_1, color=cq.Color("green")) + assy.add(cube_2, loc=cq.Location((10, 10, 10)), color=cq.Color("red")) + assy.add(cube_3, loc=cq.Location((-10, -10, -10)), color=cq.Color("red")) + assy.add(cube_4, loc=cq.Location((10, -10, -10)), color=cq.Color("red")) - # Export the assembly - success = exportStepMeta(assy, plain_step_path) - assert success + # Export the assembly, but do not use the meta STEP export method + assy.export(plain_step_path) - # Import the STEP file back in + # # Import the STEP file back in imported_assy = cq.Assembly.importStep(plain_step_path) assert imported_assy.name == "top_level" + # Check the locations + assert imported_assy.children[0].loc.toTuple()[0] == (0.0, 0.0, 0.0) + assert imported_assy.children[1].loc.toTuple()[0] == (10.0, 10.0, 10.0) + assert imported_assy.children[2].loc.toTuple()[0] == (-10.0, -10.0, -10.0) + assert imported_assy.children[3].loc.toTuple()[0] == (10.0, -10.0, -10.0) + + # Check the colors + assert pytest.approx(imported_assy.children[0].color.toTuple(), rel=0.01) == ( + 0.0, + 1.0, + 0.0, + 1.0, + ) # green + assert pytest.approx(imported_assy.children[1].color.toTuple(), rel=0.01) == ( + 1.0, + 0.0, + 0.0, + 1.0, + ) # red + assert pytest.approx(imported_assy.children[2].color.toTuple(), rel=0.01) == ( + 1.0, + 0.0, + 0.0, + 1.0, + ) # red + assert pytest.approx(imported_assy.children[3].color.toTuple(), rel=0.01) == ( + 1.0, + 0.0, + 0.0, + 1.0, + ) # red + @pytest.mark.parametrize( "assy_fixture, expected", From 43e2e00648ed182583e8d5389f28df9ab7eeb330 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Mon, 23 Jun 2025 13:54:30 -0400 Subject: [PATCH 28/44] Fixed importer for Assembly.export method --- cadquery/occ_impl/importers/assembly.py | 53 ++++++++++++++----------- tests/test_assembly.py | 42 +++++++++++++++++++- 2 files changed, 69 insertions(+), 26 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index f03384b78..aa75e75a1 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -50,7 +50,7 @@ def importStep(assy: AssemblyProtocol, path: str): color_tool = XCAFDoc_DocumentTool.ColorTool_s(doc.Main()) layer_tool = XCAFDoc_DocumentTool.LayerTool_s(doc.Main()) - def _process_simple_shape(label, parent_location=None): + def _process_simple_shape(label, parent_location=None, parent_name=None): shape = shape_tool.GetShape_s(label) # Tracks the RGB color value and whether or not it was found @@ -58,9 +58,12 @@ def _process_simple_shape(label, parent_location=None): # Load the name of the part in the assembly, if it is present name = None - name_attr = TDataStd_Name() - label.FindAttribute(TDataStd_Name.GetID_s(), name_attr) - name = str(name_attr.Get().ToExtString()) + if parent_name is not None and parent_name != assy.name: + name = parent_name + else: + name_attr = TDataStd_Name() + label.FindAttribute(TDataStd_Name.GetID_s(), name_attr) + name = str(name_attr.Get().ToExtString()) # Process the color for the shape, which could be of different types color = Quantity_Color() @@ -136,7 +139,7 @@ def _process_simple_shape(label, parent_location=None): attr_iterator.Next() - def process_label(label, parent_location=None): + def process_label(label, parent_location=None, parent_name=None): """ Recursive function that allows us to process the hierarchy of the assembly as represented in the step file. @@ -146,17 +149,16 @@ def process_label(label, parent_location=None): if shape_tool.IsReference_s(label): ref_label = TDF_Label() shape_tool.GetReferredShape_s(label, ref_label) - process_label(ref_label, parent_location) - - # Load the top-level name of the assembly, if it is present - name_attr = TDataStd_Name() - labels.Value(1).FindAttribute(TDataStd_Name.GetID_s(), name_attr) - assy.name = str(name_attr.Get().ToExtString()) + process_label(ref_label, parent_location, parent_name) return # See if this is an assembly (or sub-assembly) if shape_tool.IsAssembly_s(label): + name_attr = TDataStd_Name() + label.FindAttribute(TDataStd_Name.GetID_s(), name_attr) + name = name_attr.Get().ToExtString() + # Recursively process its components (children) comp_labels = TDF_LabelSequence() shape_tool.GetComponents_s(label, comp_labels) @@ -170,25 +172,28 @@ def process_label(label, parent_location=None): # Add the parent location if it exists if parent_location is not None: loc = parent_location * loc - process_label(sub_label, loc) + + process_label(sub_label, loc, name) + return # Check to see if we have an endpoint shape and process it if shape_tool.IsSimpleShape_s(label): - _process_simple_shape(label, parent_location) + _process_simple_shape(label, parent_location, parent_name) - # Look for the top-level assembly - found_top_level_assembly = False + # Get the shapes in the assembly labels = TDF_LabelSequence() shape_tool.GetFreeShapes(labels) - for i in range(labels.Length()): - # Make sure that we have an assembly at the top level - if shape_tool.IsTopLevel(labels.Value(i + 1)): - if shape_tool.IsAssembly_s(labels.Value(i + 1)): - found_top_level_assembly = True - - process_label(labels.Value(i + 1)) - # If we did not find a top-level assembly, raise an error - if not found_top_level_assembly: + # Use the first label to pull the top-level assembly information + name_attr = TDataStd_Name() + labels.Value(1).FindAttribute(TDataStd_Name.GetID_s(), name_attr) + assy.name = str(name_attr.Get().ToExtString()) + + # Make sure there is a top-level assembly + if shape_tool.IsTopLevel(labels.Value(1)) and shape_tool.IsAssembly_s( + labels.Value(1) + ): + process_label(labels.Value(1)) + else: raise ValueError("Step file does not contain an assembly") diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 853d5b6c0..5f8a74179 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -987,8 +987,6 @@ def test_plain_assembly_import(tmp_path_factory): Test to make sure that importing plain assemblies has not been broken. """ - from cadquery.func import box, rect - tmpdir = tmp_path_factory.mktemp("out") plain_step_path = os.path.join(tmpdir, "plain_assembly_step.step") @@ -1044,6 +1042,46 @@ def test_plain_assembly_import(tmp_path_factory): ) # red +def test_copied_assembly_import(tmp_path_factory): + """ + Tests to make sure that copied children in assemblies work correctly. + """ + from cadquery import Assembly, Location, Color + from cadquery.func import box, rect + + # Create the temporary directory + tmpdir = tmp_path_factory.mktemp("out") + + # prepare the model + def make_model(name: str, COPY: bool): + name = os.path.join(tmpdir, name) + + b = box(1, 1, 1) + + assy = Assembly(name="test_assy") + assy.add(box(1, 2, 5), color=Color("green")) + + for v in rect(10, 10).vertices(): + assy.add( + b.copy() if COPY else b, loc=Location(v.Center()), color=Color("red") + ) + + assy.export(name) + + return assy + + make_model("test_assy_copy.step", True) + make_model("test_assy.step", False) + + # import the assy with copies + assy_copy = Assembly.importStep(os.path.join(tmpdir, "test_assy_copy.step")) + assert 5 == len(assy_copy.children) + + # import the assy without copies - this throws + assy_normal = Assembly.importStep(os.path.join(tmpdir, "test_assy.step")) + assert 5 == len(assy_normal.children) + + @pytest.mark.parametrize( "assy_fixture, expected", [ From 26a64f3b4b5cb190e1e4e8edf551ab57b7437b51 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Mon, 23 Jun 2025 16:08:29 -0400 Subject: [PATCH 29/44] Fixed comment --- tests/test_assembly.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 5f8a74179..860d79cc2 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -1077,7 +1077,7 @@ def make_model(name: str, COPY: bool): assy_copy = Assembly.importStep(os.path.join(tmpdir, "test_assy_copy.step")) assert 5 == len(assy_copy.children) - # import the assy without copies - this throws + # import the assy without copies assy_normal = Assembly.importStep(os.path.join(tmpdir, "test_assy.step")) assert 5 == len(assy_normal.children) From c85ec9cd114e2a7cd2ee9aee397fa56c827ca4ea Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Mon, 30 Jun 2025 11:26:57 -0400 Subject: [PATCH 30/44] Removed a stray import that was probably added by AI somewhere along the line. --- cadquery/occ_impl/importers/assembly.py | 1 - 1 file changed, 1 deletion(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index aa75e75a1..e66f96e5d 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -14,7 +14,6 @@ import cadquery as cq from ..assembly import AssemblyProtocol -from tkinter.constants import CURRENT def importStep(assy: AssemblyProtocol, path: str): From 46c84c643ba2276766aa10d80190a69614ad0b2f Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Wed, 2 Jul 2025 15:42:07 -0400 Subject: [PATCH 31/44] Implement some of the suggestions --- cadquery/assembly.py | 4 ++-- cadquery/occ_impl/importers/assembly.py | 6 ++---- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/cadquery/assembly.py b/cadquery/assembly.py index a7bc545f1..401bac658 100644 --- a/cadquery/assembly.py +++ b/cadquery/assembly.py @@ -34,7 +34,7 @@ exportGLTF, STEPExportModeLiterals, ) -from .occ_impl.importers.assembly import importStep as importStepTopLevel +from .occ_impl.importers.assembly import importStep as _importStep from .selectors import _expression_grammar as _selector_grammar from .utils import deprecate @@ -619,7 +619,7 @@ def importStep(cls, path: str) -> Self: """ assy = cls() - importStepTopLevel(assy, path) + _importStep(assy, path) return assy diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index e66f96e5d..380ff98ce 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -52,11 +52,7 @@ def importStep(assy: AssemblyProtocol, path: str): def _process_simple_shape(label, parent_location=None, parent_name=None): shape = shape_tool.GetShape_s(label) - # Tracks the RGB color value and whether or not it was found - cq_color = None - # Load the name of the part in the assembly, if it is present - name = None if parent_name is not None and parent_name != assy.name: name = parent_name else: @@ -92,6 +88,8 @@ def _process_simple_shape(label, parent_location=None, parent_name=None): current_attr = attr_iterator.Value() # Get the type name of the attribute so that we can decide how to handle it + # TNaming_NamedShape is used to store and manage references to topological shapes, and its attributes can be accessed directly. + # XCAFDoc_GraphNode contains a graph of labels, and so we must follow the branch back to a father. if current_attr.DynamicType().Name() == "TNaming_NamedShape": # Save the shape so that we can add it to the subshape data cur_shape = current_attr.Get() From c75d0c89fd8fd56177801a8ba54b1603588f7241 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Fri, 4 Jul 2025 12:45:39 -0400 Subject: [PATCH 32/44] Added a test for nested subassemblies on import --- tests/test_assembly.py | 34 +++++++++++++++++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 860d79cc2..f8dabfd85 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -1005,7 +1005,7 @@ def test_plain_assembly_import(tmp_path_factory): # Export the assembly, but do not use the meta STEP export method assy.export(plain_step_path) - # # Import the STEP file back in + # Import the STEP file back in imported_assy = cq.Assembly.importStep(plain_step_path) assert imported_assy.name == "top_level" @@ -1082,6 +1082,38 @@ def make_model(name: str, COPY: bool): assert 5 == len(assy_normal.children) +def test_nested_subassembly_step_import(tmp_path_factory): + """ + Tests if the STEP import works correctly with nested subassemblies. + """ + + tmpdir = tmp_path_factory.mktemp("out") + nested_step_path = os.path.join(tmpdir, "plain_assembly_step.step") + + # Create a simple assembly + assy = cq.Assembly() + assy.add(cq.Workplane().box(10, 10, 10), name="box_1") + + # Create a simple subassembly + subassy = cq.Assembly() + subassy.add(cq.Workplane().box(5, 5, 5), name="box_2", loc=cq.Location(10, 10, 10)) + + # Nest the subassembly + assy.add(subassy) + + # Export and then re-import the nested assembly STEP + assy.export(nested_step_path) + imported_assy = cq.Assembly.importStep(nested_step_path) + + # Check the locations + assert imported_assy.children[0].loc.toTuple()[0] == (0.0, 0.0, 0.0) + assert imported_assy.children[1].objects["box_2"].loc.toTuple()[0] == ( + 10.0, + 10.0, + 10.0, + ) + + @pytest.mark.parametrize( "assy_fixture, expected", [ From d69856a397281b3481ec4f4b3f87bf772b698675 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Fri, 25 Jul 2025 14:50:34 -0400 Subject: [PATCH 33/44] Rework which covers everything except subshapes and layers --- cadquery/occ_impl/importers/assembly.py | 244 ++++++++++-------------- tests/test_assembly.py | 40 ++-- 2 files changed, 125 insertions(+), 159 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 380ff98ce..4899c39b0 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -1,16 +1,12 @@ from OCP.TCollection import TCollection_ExtendedString -from OCP.Quantity import Quantity_Color, Quantity_ColorRGBA -from OCP.TDocStd import TDocStd_Document +from OCP.Quantity import Quantity_ColorRGBA +from OCP.TDF import TDF_Label, TDF_LabelSequence from OCP.IFSelect import IFSelect_RetDone -from OCP.STEPCAFControl import STEPCAFControl_Reader -from OCP.XCAFDoc import ( - XCAFDoc_DocumentTool, - XCAFDoc_ColorGen, - XCAFDoc_ColorSurf, - XCAFDoc_GraphNode, -) -from OCP.TDF import TDF_Label, TDF_LabelSequence, TDF_AttributeIterator, TDF_DataSet +from OCP.TDocStd import TDocStd_Document from OCP.TDataStd import TDataStd_Name +from OCP.STEPCAFControl import STEPCAFControl_Reader +from OCP.XCAFDoc import XCAFDoc_ColorSurf +from OCP.XCAFDoc import XCAFDoc_DocumentTool import cadquery as cq from ..assembly import AssemblyProtocol @@ -26,6 +22,75 @@ def importStep(assy: AssemblyProtocol, path: str): :return: None """ + def _process_label(lbl: TDF_Label): + """ + Recursive method to process the assembly in a top-down manner. + """ + # If we have an assembly, extract all of the information out of it that we can + if shape_tool.IsAssembly_s(lbl): + # Instantiate the new assembly + new_assy = cq.Assembly() + + # Look for components + comp_labels = TDF_LabelSequence() + shape_tool.GetComponents_s(lbl, comp_labels) + + for i in range(comp_labels.Length()): + comp_label = comp_labels.Value(i + 1) + + # Get the location of the component label + loc = shape_tool.GetLocation_s(comp_label) + cq_loc = cq.Location(loc) if loc else None + + if shape_tool.IsReference_s(comp_label): + ref_label = TDF_Label() + shape_tool.GetReferredShape_s(comp_label, ref_label) + + # Find the name of this referenced part + ref_name_attr = TDataStd_Name() + if ref_label.FindAttribute(TDataStd_Name.GetID_s(), ref_name_attr): + ref_name = str(ref_name_attr.Get().ToExtString()) + + if shape_tool.IsAssembly_s(ref_label): + # Recursively process subassemblies + sub_assy = _process_label(ref_label) + + # Add the appropriate attributes to the subassembly + if sub_assy: + if cq_loc: + new_assy.add(sub_assy, name=f"{ref_name}", loc=cq_loc) + else: + new_assy.add(sub_assy, name=f"{ref_name}") + elif shape_tool.IsSimpleShape_s(ref_label): + # A single shape needs to be added to the assembly + final_shape = shape_tool.GetShape_s(ref_label) + cq_shape = cq.Shape.cast(final_shape) + + # Find the subshape color, if there is one set for this shape + color = Quantity_ColorRGBA() + # Extract the color, if present on the shape + if color_tool.GetColor(final_shape, XCAFDoc_ColorSurf, color): + rgb = color.GetRGB() + cq_color = cq.Color( + rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha() + ) + else: + cq_color = None + + if cq_loc: + new_assy.add( + cq_shape, name=f"{ref_name}", loc=cq_loc, color=cq_color + ) + else: + new_assy.add(cq_shape, name=f"{ref_name}", color=cq_color) + + return new_assy + elif shape_tool.IsSimpleShape_s(lbl): + shape = shape_tool.GetShape_s(lbl) + return cq.Shape.cast(shape) + else: + return None + # Document that the step file will be read into doc = TDocStd_Document(TCollection_ExtendedString("XmlOcaf")) @@ -49,148 +114,33 @@ def importStep(assy: AssemblyProtocol, path: str): color_tool = XCAFDoc_DocumentTool.ColorTool_s(doc.Main()) layer_tool = XCAFDoc_DocumentTool.LayerTool_s(doc.Main()) - def _process_simple_shape(label, parent_location=None, parent_name=None): - shape = shape_tool.GetShape_s(label) - - # Load the name of the part in the assembly, if it is present - if parent_name is not None and parent_name != assy.name: - name = parent_name - else: - name_attr = TDataStd_Name() - label.FindAttribute(TDataStd_Name.GetID_s(), name_attr) - name = str(name_attr.Get().ToExtString()) - - # Process the color for the shape, which could be of different types - color = Quantity_Color() - cq_color = cq.Color(0.50, 0.50, 0.50) - if color_tool.GetColor_s(label, XCAFDoc_ColorSurf, color): - r = color.Red() - g = color.Green() - b = color.Blue() - cq_color = cq.Color(r, g, b) - - # Handle the location if it was passed down form a parent component - if parent_location is not None: - assy.add( - cq.Shape.cast(shape), - name=name, - color=cq_color, - loc=cq.Location(parent_location), - ) - else: - assy.add(cq.Shape.cast(shape), name=name, color=cq_color) - - # Check all the attributes of all the children to find the subshapes and any names - for j in range(label.NbChildren()): - child_label = label.FindChild(j + 1) - attr_iterator = TDF_AttributeIterator(child_label) - while attr_iterator.More(): - current_attr = attr_iterator.Value() - - # Get the type name of the attribute so that we can decide how to handle it - # TNaming_NamedShape is used to store and manage references to topological shapes, and its attributes can be accessed directly. - # XCAFDoc_GraphNode contains a graph of labels, and so we must follow the branch back to a father. - if current_attr.DynamicType().Name() == "TNaming_NamedShape": - # Save the shape so that we can add it to the subshape data - cur_shape = current_attr.Get() - - # Find the layer name, if there is one set for this shape - layers = TDF_LabelSequence() - layer_tool.GetLayers(child_label, layers) - for i in range(1, layers.Length() + 1): - lbl = layers.Value(i) - name_attr = TDataStd_Name() - lbl.FindAttribute(TDataStd_Name.GetID_s(), name_attr) - - # Extract the layer name for the shape here - layer_name = name_attr.Get().ToExtString() - - # Add the layer as a subshape entry on the assembly - assy.addSubshape(cur_shape, layer=layer_name) - - # Find the subshape color, if there is one set for this shape - color = Quantity_ColorRGBA() - # Extract the color, if present on the shape - if color_tool.GetColor(cur_shape, XCAFDoc_ColorSurf, color): - rgb = color.GetRGB() - cq_color = cq.Color( - rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha() - ) - - # Save the color info via the assembly subshape mechanism - assy.addSubshape(cur_shape, color=cq_color) - elif current_attr.DynamicType().Name() == "XCAFDoc_GraphNode": - # Step up one level to try to get the name from the parent - lbl = current_attr.GetFather(1).Label() - - # Step through and search for the name attribute - it = TDF_AttributeIterator(lbl) - while it.More(): - new_attr = it.Value() - if new_attr.DynamicType().Name() == "TDataStd_Name": - # Save this as the name of the subshape - assy.addSubshape( - cur_shape, name=new_attr.Get().ToExtString(), - ) - it.Next() - - attr_iterator.Next() - - def process_label(label, parent_location=None, parent_name=None): - """ - Recursive function that allows us to process the hierarchy of the assembly as represented - in the step file. - """ - - # Handle reference labels - if shape_tool.IsReference_s(label): - ref_label = TDF_Label() - shape_tool.GetReferredShape_s(label, ref_label) - process_label(ref_label, parent_location, parent_name) - - return - - # See if this is an assembly (or sub-assembly) - if shape_tool.IsAssembly_s(label): - name_attr = TDataStd_Name() - label.FindAttribute(TDataStd_Name.GetID_s(), name_attr) - name = name_attr.Get().ToExtString() - - # Recursively process its components (children) - comp_labels = TDF_LabelSequence() - shape_tool.GetComponents_s(label, comp_labels) - - for i in range(comp_labels.Length()): - sub_label = comp_labels.Value(i + 1) - # Get the location of the sub-label, if it exists - loc = shape_tool.GetLocation_s(sub_label) - - # Pass down the location or other context as needed - # Add the parent location if it exists - if parent_location is not None: - loc = parent_location * loc - - process_label(sub_label, loc, name) - - return - - # Check to see if we have an endpoint shape and process it - if shape_tool.IsSimpleShape_s(label): - _process_simple_shape(label, parent_location, parent_name) - - # Get the shapes in the assembly + # Collect all the labels representing shapes in the document labels = TDF_LabelSequence() shape_tool.GetFreeShapes(labels) + if labels.Length() == 0: + raise ValueError("No assembly found in STEP file") - # Use the first label to pull the top-level assembly information - name_attr = TDataStd_Name() - labels.Value(1).FindAttribute(TDataStd_Name.GetID_s(), name_attr) - assy.name = str(name_attr.Get().ToExtString()) + # Get the top-level label, which should represent an assembly + top_level_label = labels.Value(1) # Make sure there is a top-level assembly - if shape_tool.IsTopLevel(labels.Value(1)) and shape_tool.IsAssembly_s( - labels.Value(1) + if shape_tool.IsTopLevel(top_level_label) and shape_tool.IsAssembly_s( + top_level_label ): - process_label(labels.Value(1)) + # Set the name of the top-level assembly to match the top-level label + name_attr = TDataStd_Name() + top_level_label.FindAttribute(TDataStd_Name.GetID_s(), name_attr) + assy.name = str(name_attr.Get().ToExtString()) + + # Start the recursive processing of labels + whole_assy = _process_label(top_level_label) + + # Copy contents instead of adding the whole assembly + # assy.name = whole_assy.name + if whole_assy and hasattr(whole_assy, "children"): + for child in whole_assy.children: + assy.add(child, name=child.name, color=child.color, loc=child.loc) + elif whole_assy: + assy.add(whole_assy) else: raise ValueError("Step file does not contain an assembly") diff --git a/tests/test_assembly.py b/tests/test_assembly.py index f8dabfd85..7a3da384b 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -1010,31 +1010,47 @@ def test_plain_assembly_import(tmp_path_factory): assert imported_assy.name == "top_level" # Check the locations - assert imported_assy.children[0].loc.toTuple()[0] == (0.0, 0.0, 0.0) - assert imported_assy.children[1].loc.toTuple()[0] == (10.0, 10.0, 10.0) - assert imported_assy.children[2].loc.toTuple()[0] == (-10.0, -10.0, -10.0) - assert imported_assy.children[3].loc.toTuple()[0] == (10.0, -10.0, -10.0) + assert imported_assy.children[0].children[0].loc.toTuple()[0] == (0.0, 0.0, 0.0,) + assert imported_assy.children[0].children[1].loc.toTuple()[0] == (10.0, 10.0, 10.0,) + assert imported_assy.children[0].children[2].loc.toTuple()[0] == ( + -10.0, + -10.0, + -10.0, + ) + assert imported_assy.children[0].children[3].loc.toTuple()[0] == ( + 10.0, + -10.0, + -10.0, + ) # Check the colors - assert pytest.approx(imported_assy.children[0].color.toTuple(), rel=0.01) == ( + assert pytest.approx( + imported_assy.children[0].children[0].children[0].color.toTuple(), rel=0.01 + ) == ( 0.0, 1.0, 0.0, 1.0, ) # green - assert pytest.approx(imported_assy.children[1].color.toTuple(), rel=0.01) == ( + assert pytest.approx( + imported_assy.children[0].children[1].children[0].color.toTuple(), rel=0.01 + ) == ( 1.0, 0.0, 0.0, 1.0, ) # red - assert pytest.approx(imported_assy.children[2].color.toTuple(), rel=0.01) == ( + assert pytest.approx( + imported_assy.children[0].children[2].children[0].color.toTuple(), rel=0.01 + ) == ( 1.0, 0.0, 0.0, 1.0, ) # red - assert pytest.approx(imported_assy.children[3].color.toTuple(), rel=0.01) == ( + assert pytest.approx( + imported_assy.children[0].children[3].children[0].color.toTuple(), rel=0.01 + ) == ( 1.0, 0.0, 0.0, @@ -1075,11 +1091,11 @@ def make_model(name: str, COPY: bool): # import the assy with copies assy_copy = Assembly.importStep(os.path.join(tmpdir, "test_assy_copy.step")) - assert 5 == len(assy_copy.children) + assert 5 == len(assy_copy.children[0].children) # import the assy without copies assy_normal = Assembly.importStep(os.path.join(tmpdir, "test_assy.step")) - assert 5 == len(assy_normal.children) + assert 5 == len(assy_normal.children[0].children) def test_nested_subassembly_step_import(tmp_path_factory): @@ -1106,8 +1122,8 @@ def test_nested_subassembly_step_import(tmp_path_factory): imported_assy = cq.Assembly.importStep(nested_step_path) # Check the locations - assert imported_assy.children[0].loc.toTuple()[0] == (0.0, 0.0, 0.0) - assert imported_assy.children[1].objects["box_2"].loc.toTuple()[0] == ( + assert imported_assy.children[0].children[0].loc.toTuple()[0] == (0.0, 0.0, 0.0) + assert imported_assy.children[0].children[1].objects["box_2"].loc.toTuple()[0] == ( 10.0, 10.0, 10.0, From c8e0c08e60bb0d3d7f9c7bec8e03cb0b92d52f35 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Fri, 25 Jul 2025 15:33:35 -0400 Subject: [PATCH 34/44] Added layer name support back in --- cadquery/occ_impl/importers/assembly.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 4899c39b0..2d8f94efe 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -84,6 +84,20 @@ def _process_label(lbl: TDF_Label): else: new_assy.add(cq_shape, name=f"{ref_name}", color=cq_color) + # Find the layer name, if there is one set for this shape + layers = TDF_LabelSequence() + layer_tool.GetLayers(ref_label, layers) + for i in range(1, layers.Length() + 1): + lbl = layers.Value(i) + name_attr = TDataStd_Name() + lbl.FindAttribute(TDataStd_Name.GetID_s(), name_attr) + + # Extract the layer name for the shape here + layer_name = name_attr.Get().ToExtString() + + # Add the layer as a subshape entry on the assembly + new_assy.addSubshape(final_shape, layer=layer_name) + return new_assy elif shape_tool.IsSimpleShape_s(lbl): shape = shape_tool.GetShape_s(lbl) From 7639fdf4bcd3697154f8c8ba2401bf5c8db5e18d Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Mon, 28 Jul 2025 09:21:56 -0400 Subject: [PATCH 35/44] Added a round-trip test and fixed issues that it revealed --- cadquery/occ_impl/importers/assembly.py | 101 ++++++++++++++++++++---- tests/test_assembly.py | 58 ++++++++++++++ 2 files changed, 143 insertions(+), 16 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 2d8f94efe..d9a3bcc12 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -1,6 +1,6 @@ from OCP.TCollection import TCollection_ExtendedString from OCP.Quantity import Quantity_ColorRGBA -from OCP.TDF import TDF_Label, TDF_LabelSequence +from OCP.TDF import TDF_Label, TDF_LabelSequence, TDF_AttributeIterator from OCP.IFSelect import IFSelect_RetDone from OCP.TDocStd import TDocStd_Document from OCP.TDataStd import TDataStd_Name @@ -84,19 +84,82 @@ def _process_label(lbl: TDF_Label): else: new_assy.add(cq_shape, name=f"{ref_name}", color=cq_color) - # Find the layer name, if there is one set for this shape - layers = TDF_LabelSequence() - layer_tool.GetLayers(ref_label, layers) - for i in range(1, layers.Length() + 1): - lbl = layers.Value(i) - name_attr = TDataStd_Name() - lbl.FindAttribute(TDataStd_Name.GetID_s(), name_attr) - - # Extract the layer name for the shape here - layer_name = name_attr.Get().ToExtString() - - # Add the layer as a subshape entry on the assembly - new_assy.addSubshape(final_shape, layer=layer_name) + # Search for subshape names, layers and colors + for j in range(ref_label.NbChildren()): + child_label = ref_label.FindChild(j + 1) + + # Iterate through all the attributes looking for subshapes + attr_iterator = TDF_AttributeIterator(child_label) + while attr_iterator.More(): + current_attr = attr_iterator.Value() + + # TNaming_NamedShape is used to store and manage references to + # topological shapes, and its attributes can be accessed directly. + # XCAFDoc_GraphNode contains a graph of labels, and so we must + # follow the branch back to a father. + if ( + current_attr.DynamicType().Name() + == "XCAFDoc_GraphNode" + ): + # Step up one level to try to get the name from the parent + lbl = current_attr.GetFather(1).Label() + + # Step through and search for the name attribute + it = TDF_AttributeIterator(lbl) + while it.More(): + new_attr = it.Value() + if ( + new_attr.DynamicType().Name() + == "TDataStd_Name" + ): + # Save this as the name of the subshape + assy.addSubshape( + cur_shape, + name=new_attr.Get().ToExtString(), + ) + break + it.Next() + elif ( + current_attr.DynamicType().Name() + == "TNaming_NamedShape" + ): + # Save the shape so that we can add it to the subshape data + cur_shape = current_attr.Get() + + # Find the layer name, if there is one set for this shape + layers = TDF_LabelSequence() + layer_tool.GetLayers(child_label, layers) + for i in range(1, layers.Length() + 1): + lbl = layers.Value(i) + name_attr = TDataStd_Name() + lbl.FindAttribute( + TDataStd_Name.GetID_s(), name_attr + ) + + # Extract the layer name for the shape here + layer_name = name_attr.Get().ToExtString() + + # Add the layer as a subshape entry on the assembly + assy.addSubshape(cur_shape, layer=layer_name) + + # Find the subshape color, if there is one set for this shape + color = Quantity_ColorRGBA() + # Extract the color, if present on the shape + if color_tool.GetColor( + cur_shape, XCAFDoc_ColorSurf, color + ): + rgb = color.GetRGB() + cq_color = cq.Color( + rgb.Red(), + rgb.Green(), + rgb.Blue(), + color.Alpha(), + ) + + # Save the color info via the assembly subshape mechanism + assy.addSubshape(cur_shape, color=cq_color) + + attr_iterator.Next() return new_assy elif shape_tool.IsSimpleShape_s(lbl): @@ -149,9 +212,15 @@ def _process_label(lbl: TDF_Label): # Start the recursive processing of labels whole_assy = _process_label(top_level_label) - # Copy contents instead of adding the whole assembly - # assy.name = whole_assy.name if whole_assy and hasattr(whole_assy, "children"): + # Check to see if there is an extra top-level node. This is done because + # cq.Assembly.export adds an extra top-level node which will cause a cascade of + # extras on successive round-trips. exportStepMeta does not add the extra top-level + # node and so does not exhibit this behavior. + if assy.name == whole_assy.children[0].name: + whole_assy = whole_assy.children[0] + + # Copy all of the children over to the main assembly object for child in whole_assy.children: assy.add(child, name=child.name, color=child.color, loc=child.loc) elif whole_assy: diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 7a3da384b..e251db332 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -1130,6 +1130,64 @@ def test_nested_subassembly_step_import(tmp_path_factory): ) +def test_assembly_step_import_roundtrip(tmp_path_factory): + """ + Tests that the assembly does not mutate during successive export-import round trips. + """ + + # Set up the temporary directory + tmpdir = tmp_path_factory.mktemp("out") + round_trip_step_path = os.path.join(tmpdir, "round_trip.step") + + # Create a sample assembly + assy = cq.Assembly(name="top-level") + assy.add(cq.Workplane().box(10, 10, 10), name="cube_1", color=cq.Color("red")) + subshape_assy = cq.Assembly(name="nested-assy") + subshape_assy.add( + cq.Workplane().cylinder(height=10.0, radius=2.5), + name="cylinder_1", + color=cq.Color("blue"), + loc=cq.Location((20, 20, 20)), + ) + assy.add(subshape_assy) + + # First export + assy.export(round_trip_step_path) + + # First import + assy = cq.Assembly.importStep(round_trip_step_path) + + # Second export + assy.export(round_trip_step_path) + + # Second import + assy = cq.Assembly.importStep(round_trip_step_path) + + # Check some general aspects of the assembly structure now + assert len(assy.children) == 2 + assert assy.name == "top-level" + assert assy.children[0].name == "cube_1" + assert assy.children[1].children[0].name == "cylinder_1" + + # First meta export + exportStepMeta(assy, round_trip_step_path) + + # First meta import + assy = cq.Assembly.importStep(round_trip_step_path) + + # Second meta export + exportStepMeta(assy, round_trip_step_path) + + # Second meta import + assy = cq.Assembly.importStep(round_trip_step_path) + + # Check some general aspects of the assembly structure now + assert len(assy.children) == 2 + assert assy.name == "top-level" + assert assy.children[0].name == "cube_1" + assert assy.children[1].children[0].name == "cylinder_1" + + @pytest.mark.parametrize( "assy_fixture, expected", [ From eedc16f0db1d66146017c852613464fe9b50cd0b Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Mon, 28 Jul 2025 15:55:31 -0400 Subject: [PATCH 36/44] mypy fixes --- cadquery/occ_impl/assembly.py | 33 +++++++++++++++++++++++++ cadquery/occ_impl/importers/assembly.py | 33 ++++++++++++++++--------- tests/test_assembly.py | 32 +++++++++--------------- 3 files changed, 67 insertions(+), 31 deletions(-) diff --git a/cadquery/occ_impl/assembly.py b/cadquery/occ_impl/assembly.py index d06282e53..4cc30f8e0 100644 --- a/cadquery/occ_impl/assembly.py +++ b/cadquery/occ_impl/assembly.py @@ -184,6 +184,39 @@ def _subshape_colors(self) -> Dict[Shape, Color]: def _subshape_layers(self) -> Dict[Shape, str]: ... + @overload + def add( + self, + obj: "Assembly", + loc: Optional[Location] = None, + name: Optional[str] = None, + color: Optional[Color] = None, + ): + ... + + @overload + def add( + self, + obj: AssemblyObjects, + loc: Optional[Location] = None, + name: Optional[str] = None, + color: Optional[Color] = None, + metadata: Optional[Dict[str, Any]] = None, + ): + ... + + def add(self, arg, **kwargs): + ... + + def addSubshape( + self, + s: Shape, + name: Optional[str] = None, + color: Optional[Color] = None, + layer: Optional[str] = None, + ) -> "Assembly": + ... + def traverse(self) -> Iterable[Tuple[str, "AssemblyProtocol"]]: ... diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index d9a3bcc12..0d4f001a6 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -1,3 +1,4 @@ +from OCP.TopoDS import TopoDS_Shape from OCP.TCollection import TCollection_ExtendedString from OCP.Quantity import Quantity_ColorRGBA from OCP.TDF import TDF_Label, TDF_LabelSequence, TDF_AttributeIterator @@ -113,7 +114,7 @@ def _process_label(lbl: TDF_Label): == "TDataStd_Name" ): # Save this as the name of the subshape - assy.addSubshape( + new_assy.addSubshape( cur_shape, name=new_attr.Get().ToExtString(), ) @@ -124,7 +125,7 @@ def _process_label(lbl: TDF_Label): == "TNaming_NamedShape" ): # Save the shape so that we can add it to the subshape data - cur_shape = current_attr.Get() + cur_shape: TopoDS_Shape = current_attr.Get() # Find the layer name, if there is one set for this shape layers = TDF_LabelSequence() @@ -140,7 +141,9 @@ def _process_label(lbl: TDF_Label): layer_name = name_attr.Get().ToExtString() # Add the layer as a subshape entry on the assembly - assy.addSubshape(cur_shape, layer=layer_name) + new_assy.addSubshape( + cur_shape, layer=layer_name + ) # Find the subshape color, if there is one set for this shape color = Quantity_ColorRGBA() @@ -157,7 +160,7 @@ def _process_label(lbl: TDF_Label): ) # Save the color info via the assembly subshape mechanism - assy.addSubshape(cur_shape, color=cq_color) + new_assy.addSubshape(cur_shape, color=cq_color) attr_iterator.Next() @@ -210,20 +213,28 @@ def _process_label(lbl: TDF_Label): assy.name = str(name_attr.Get().ToExtString()) # Start the recursive processing of labels - whole_assy = _process_label(top_level_label) + imported_assy = _process_label(top_level_label) - if whole_assy and hasattr(whole_assy, "children"): + if imported_assy and hasattr(imported_assy, "children"): # Check to see if there is an extra top-level node. This is done because # cq.Assembly.export adds an extra top-level node which will cause a cascade of # extras on successive round-trips. exportStepMeta does not add the extra top-level # node and so does not exhibit this behavior. - if assy.name == whole_assy.children[0].name: - whole_assy = whole_assy.children[0] + if assy.name == imported_assy.children[0].name: + imported_assy = imported_assy.children[0] # Copy all of the children over to the main assembly object - for child in whole_assy.children: + for child in imported_assy.children: assy.add(child, name=child.name, color=child.color, loc=child.loc) - elif whole_assy: - assy.add(whole_assy) + elif imported_assy: + assy.add(imported_assy) + + # Copy across subshape data + for shape, name in imported_assy._subshape_names.items(): + assy.addSubshape(shape, name=name) + for shape, color in imported_assy._subshape_colors.items(): + assy.addSubshape(shape, color=color) + for shape, layer in imported_assy._subshape_layers.items(): + assy.addSubshape(shape, layer=layer) else: raise ValueError("Step file does not contain an assembly") diff --git a/tests/test_assembly.py b/tests/test_assembly.py index e251db332..f045671d5 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -1010,22 +1010,14 @@ def test_plain_assembly_import(tmp_path_factory): assert imported_assy.name == "top_level" # Check the locations - assert imported_assy.children[0].children[0].loc.toTuple()[0] == (0.0, 0.0, 0.0,) - assert imported_assy.children[0].children[1].loc.toTuple()[0] == (10.0, 10.0, 10.0,) - assert imported_assy.children[0].children[2].loc.toTuple()[0] == ( - -10.0, - -10.0, - -10.0, - ) - assert imported_assy.children[0].children[3].loc.toTuple()[0] == ( - 10.0, - -10.0, - -10.0, - ) + assert imported_assy.children[0].loc.toTuple()[0] == (0.0, 0.0, 0.0,) + assert imported_assy.children[1].loc.toTuple()[0] == (10.0, 10.0, 10.0,) + assert imported_assy.children[2].loc.toTuple()[0] == (-10.0, -10.0, -10.0,) + assert imported_assy.children[3].loc.toTuple()[0] == (10.0, -10.0, -10.0,) # Check the colors assert pytest.approx( - imported_assy.children[0].children[0].children[0].color.toTuple(), rel=0.01 + imported_assy.children[0].children[0].color.toTuple(), rel=0.01 ) == ( 0.0, 1.0, @@ -1033,7 +1025,7 @@ def test_plain_assembly_import(tmp_path_factory): 1.0, ) # green assert pytest.approx( - imported_assy.children[0].children[1].children[0].color.toTuple(), rel=0.01 + imported_assy.children[1].children[0].color.toTuple(), rel=0.01 ) == ( 1.0, 0.0, @@ -1041,7 +1033,7 @@ def test_plain_assembly_import(tmp_path_factory): 1.0, ) # red assert pytest.approx( - imported_assy.children[0].children[2].children[0].color.toTuple(), rel=0.01 + imported_assy.children[2].children[0].color.toTuple(), rel=0.01 ) == ( 1.0, 0.0, @@ -1049,7 +1041,7 @@ def test_plain_assembly_import(tmp_path_factory): 1.0, ) # red assert pytest.approx( - imported_assy.children[0].children[3].children[0].color.toTuple(), rel=0.01 + imported_assy.children[3].children[0].color.toTuple(), rel=0.01 ) == ( 1.0, 0.0, @@ -1091,11 +1083,11 @@ def make_model(name: str, COPY: bool): # import the assy with copies assy_copy = Assembly.importStep(os.path.join(tmpdir, "test_assy_copy.step")) - assert 5 == len(assy_copy.children[0].children) + assert 5 == len(assy_copy.children) # import the assy without copies assy_normal = Assembly.importStep(os.path.join(tmpdir, "test_assy.step")) - assert 5 == len(assy_normal.children[0].children) + assert 5 == len(assy_normal.children) def test_nested_subassembly_step_import(tmp_path_factory): @@ -1122,8 +1114,8 @@ def test_nested_subassembly_step_import(tmp_path_factory): imported_assy = cq.Assembly.importStep(nested_step_path) # Check the locations - assert imported_assy.children[0].children[0].loc.toTuple()[0] == (0.0, 0.0, 0.0) - assert imported_assy.children[0].children[1].objects["box_2"].loc.toTuple()[0] == ( + assert imported_assy.children[0].loc.toTuple()[0] == (0.0, 0.0, 0.0) + assert imported_assy.children[1].objects["box_2"].loc.toTuple()[0] == ( 10.0, 10.0, 10.0, From ee2b83475704c02409952e5af1c6c0693a95c410 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Mon, 28 Jul 2025 16:40:00 -0400 Subject: [PATCH 37/44] More mypy fixes --- cadquery/occ_impl/assembly.py | 4 ++-- cadquery/occ_impl/importers/assembly.py | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/cadquery/occ_impl/assembly.py b/cadquery/occ_impl/assembly.py index 4cc30f8e0..9f5b1af8c 100644 --- a/cadquery/occ_impl/assembly.py +++ b/cadquery/occ_impl/assembly.py @@ -187,7 +187,7 @@ def _subshape_layers(self) -> Dict[Shape, str]: @overload def add( self, - obj: "Assembly", + obj: "AssemblyProtocol", loc: Optional[Location] = None, name: Optional[str] = None, color: Optional[Color] = None, @@ -214,7 +214,7 @@ def addSubshape( name: Optional[str] = None, color: Optional[Color] = None, layer: Optional[str] = None, - ) -> "Assembly": + ) -> "AssemblyProtocol": ... def traverse(self) -> Iterable[Tuple[str, "AssemblyProtocol"]]: diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 0d4f001a6..84d587ad0 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -142,7 +142,7 @@ def _process_label(lbl: TDF_Label): # Add the layer as a subshape entry on the assembly new_assy.addSubshape( - cur_shape, layer=layer_name + cq.Shape.cast(cur_shape), layer=layer_name ) # Find the subshape color, if there is one set for this shape @@ -160,7 +160,9 @@ def _process_label(lbl: TDF_Label): ) # Save the color info via the assembly subshape mechanism - new_assy.addSubshape(cur_shape, color=cq_color) + new_assy.addSubshape( + cq.Shape.cast(cur_shape), color=cq_color + ) attr_iterator.Next() From 2b8a8eb61437b690fd146436ad32803361d8e607 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Mon, 28 Jul 2025 17:00:50 -0400 Subject: [PATCH 38/44] Missed a cast --- cadquery/occ_impl/importers/assembly.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 84d587ad0..a390d5198 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -115,7 +115,7 @@ def _process_label(lbl: TDF_Label): ): # Save this as the name of the subshape new_assy.addSubshape( - cur_shape, + cq.Shape.cast(cur_shape), name=new_attr.Get().ToExtString(), ) break From c7821ee2f3f993f7c32d3d2a9ea032610dcba2ee Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Tue, 29 Jul 2025 17:47:47 -0400 Subject: [PATCH 39/44] More mypy fixes --- cadquery/occ_impl/assembly.py | 23 +++++++++++---- cadquery/occ_impl/importers/assembly.py | 38 ++++++++++++------------- 2 files changed, 36 insertions(+), 25 deletions(-) diff --git a/cadquery/occ_impl/assembly.py b/cadquery/occ_impl/assembly.py index 9f5b1af8c..4b00bb1d0 100644 --- a/cadquery/occ_impl/assembly.py +++ b/cadquery/occ_impl/assembly.py @@ -10,7 +10,7 @@ List, cast, ) -from typing_extensions import Protocol +from typing_extensions import Protocol, Self from math import degrees, radians from OCP.TDocStd import TDocStd_Document @@ -187,11 +187,11 @@ def _subshape_layers(self) -> Dict[Shape, str]: @overload def add( self, - obj: "AssemblyProtocol", + obj: Self, loc: Optional[Location] = None, name: Optional[str] = None, color: Optional[Color] = None, - ): + ) -> Self: ... @overload @@ -202,10 +202,21 @@ def add( name: Optional[str] = None, color: Optional[Color] = None, metadata: Optional[Dict[str, Any]] = None, - ): + ) -> Self: ... - def add(self, arg, **kwargs): + def add( + self, + obj: Union[Self, AssemblyObjects], + loc: Optional[Location] = None, + name: Optional[str] = None, + color: Optional[Color] = None, + metadata: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> Self: + """ + Add a subassembly to the current assembly. + """ ... def addSubshape( @@ -214,7 +225,7 @@ def addSubshape( name: Optional[str] = None, color: Optional[Color] = None, layer: Optional[str] = None, - ) -> "AssemblyProtocol": + ) -> Self: ... def traverse(self) -> Iterable[Tuple[str, "AssemblyProtocol"]]: diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index a390d5198..c88e132a4 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -102,30 +102,30 @@ def _process_label(lbl: TDF_Label): current_attr.DynamicType().Name() == "XCAFDoc_GraphNode" ): - # Step up one level to try to get the name from the parent - lbl = current_attr.GetFather(1).Label() - - # Step through and search for the name attribute - it = TDF_AttributeIterator(lbl) - while it.More(): - new_attr = it.Value() - if ( - new_attr.DynamicType().Name() - == "TDataStd_Name" - ): - # Save this as the name of the subshape - new_assy.addSubshape( - cq.Shape.cast(cur_shape), - name=new_attr.Get().ToExtString(), - ) - break - it.Next() + # Compatibility check + if hasattr(current_attr, "GetFather"): + lbl = current_attr.GetFather(1).Label() + else: + lbl = current_attr.Label().Father() + + # Find the name attribute and add it for the subshape + name_attr = TDataStd_Name() + if lbl.FindAttribute( + TDataStd_Name.GetID_s(), name_attr + ): + # Save this as the name of the subshape + new_assy.addSubshape( + cq.Shape.cast(cur_shape), + name=name_attr.Get().ToExtString(), + ) elif ( current_attr.DynamicType().Name() == "TNaming_NamedShape" ): # Save the shape so that we can add it to the subshape data - cur_shape: TopoDS_Shape = current_attr.Get() + cur_shape: TopoDS_Shape = shape_tool.GetShape_s( + child_label + ) # Find the layer name, if there is one set for this shape layers = TDF_LabelSequence() From f88aa9a57fdbcda585e3fac7c7e7f8e4b9e865dc Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Wed, 30 Jul 2025 12:32:59 -0400 Subject: [PATCH 40/44] Tried to remove the attribute iterator and could not, but moved some code out of loop --- cadquery/occ_impl/importers/assembly.py | 80 +++++++++++++------------ 1 file changed, 42 insertions(+), 38 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index c88e132a4..ed4e5d98b 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -89,7 +89,48 @@ def _process_label(lbl: TDF_Label): for j in range(ref_label.NbChildren()): child_label = ref_label.FindChild(j + 1) - # Iterate through all the attributes looking for subshapes + # Save the shape so that we can add it to the subshape data + cur_shape: TopoDS_Shape = shape_tool.GetShape_s(child_label) + if cur_shape.IsNull(): + continue + + # Validate the child label before using it + if not child_label or child_label.IsNull(): + continue + + # Find the layer name, if there is one set for this shape + layers = TDF_LabelSequence() + layer_tool.GetLayers(child_label, layers) + for i in range(1, layers.Length() + 1): + lbl = layers.Value(i) + name_attr = TDataStd_Name() + lbl.FindAttribute(TDataStd_Name.GetID_s(), name_attr) + + # Extract the layer name for the shape here + layer_name = name_attr.Get().ToExtString() + + # Add the layer as a subshape entry on the assembly + new_assy.addSubshape( + cq.Shape.cast(cur_shape), layer=layer_name + ) + + # Find the subshape color, if there is one set for this shape + color = Quantity_ColorRGBA() + # Extract the color, if present on the shape + if color_tool.GetColor(cur_shape, XCAFDoc_ColorSurf, color): + rgb = color.GetRGB() + cq_color = cq.Color( + rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha(), + ) + + # Save the color info via the assembly subshape mechanism + new_assy.addSubshape( + cq.Shape.cast(cur_shape), color=cq_color + ) + + # Iterate through all the attributes looking for subshape names. + # This is safer than trying to access the attributes directly with + # FindAttribute because it will cause a segfault in certain cases. attr_iterator = TDF_AttributeIterator(child_label) while attr_iterator.More(): current_attr = attr_iterator.Value() @@ -127,43 +168,6 @@ def _process_label(lbl: TDF_Label): child_label ) - # Find the layer name, if there is one set for this shape - layers = TDF_LabelSequence() - layer_tool.GetLayers(child_label, layers) - for i in range(1, layers.Length() + 1): - lbl = layers.Value(i) - name_attr = TDataStd_Name() - lbl.FindAttribute( - TDataStd_Name.GetID_s(), name_attr - ) - - # Extract the layer name for the shape here - layer_name = name_attr.Get().ToExtString() - - # Add the layer as a subshape entry on the assembly - new_assy.addSubshape( - cq.Shape.cast(cur_shape), layer=layer_name - ) - - # Find the subshape color, if there is one set for this shape - color = Quantity_ColorRGBA() - # Extract the color, if present on the shape - if color_tool.GetColor( - cur_shape, XCAFDoc_ColorSurf, color - ): - rgb = color.GetRGB() - cq_color = cq.Color( - rgb.Red(), - rgb.Green(), - rgb.Blue(), - color.Alpha(), - ) - - # Save the color info via the assembly subshape mechanism - new_assy.addSubshape( - cq.Shape.cast(cur_shape), color=cq_color - ) - attr_iterator.Next() return new_assy From 71e309ab2df95a6aeacd4aac109c4cc66ec9ace2 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Wed, 30 Jul 2025 13:05:38 -0400 Subject: [PATCH 41/44] Fixes and simplifications based on codecov checks --- cadquery/occ_impl/importers/assembly.py | 56 +++++++------------------ 1 file changed, 15 insertions(+), 41 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index ed4e5d98b..62b4b212c 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -57,11 +57,7 @@ def _process_label(lbl: TDF_Label): sub_assy = _process_label(ref_label) # Add the appropriate attributes to the subassembly - if sub_assy: - if cq_loc: - new_assy.add(sub_assy, name=f"{ref_name}", loc=cq_loc) - else: - new_assy.add(sub_assy, name=f"{ref_name}") + new_assy.add(sub_assy, name=f"{ref_name}", loc=cq_loc) elif shape_tool.IsSimpleShape_s(ref_label): # A single shape needs to be added to the assembly final_shape = shape_tool.GetShape_s(ref_label) @@ -78,12 +74,9 @@ def _process_label(lbl: TDF_Label): else: cq_color = None - if cq_loc: - new_assy.add( - cq_shape, name=f"{ref_name}", loc=cq_loc, color=cq_color - ) - else: - new_assy.add(cq_shape, name=f"{ref_name}", color=cq_color) + new_assy.add( + cq_shape, name=f"{ref_name}", loc=cq_loc, color=cq_color + ) # Search for subshape names, layers and colors for j in range(ref_label.NbChildren()): @@ -91,12 +84,6 @@ def _process_label(lbl: TDF_Label): # Save the shape so that we can add it to the subshape data cur_shape: TopoDS_Shape = shape_tool.GetShape_s(child_label) - if cur_shape.IsNull(): - continue - - # Validate the child label before using it - if not child_label or child_label.IsNull(): - continue # Find the layer name, if there is one set for this shape layers = TDF_LabelSequence() @@ -143,11 +130,9 @@ def _process_label(lbl: TDF_Label): current_attr.DynamicType().Name() == "XCAFDoc_GraphNode" ): - # Compatibility check + # Only the GraphNode should have this method if hasattr(current_attr, "GetFather"): lbl = current_attr.GetFather(1).Label() - else: - lbl = current_attr.Label().Father() # Find the name attribute and add it for the subshape name_attr = TDataStd_Name() @@ -171,11 +156,6 @@ def _process_label(lbl: TDF_Label): attr_iterator.Next() return new_assy - elif shape_tool.IsSimpleShape_s(lbl): - shape = shape_tool.GetShape_s(lbl) - return cq.Shape.cast(shape) - else: - return None # Document that the step file will be read into doc = TDocStd_Document(TCollection_ExtendedString("XmlOcaf")) @@ -203,12 +183,9 @@ def _process_label(lbl: TDF_Label): # Collect all the labels representing shapes in the document labels = TDF_LabelSequence() shape_tool.GetFreeShapes(labels) - if labels.Length() == 0: - raise ValueError("No assembly found in STEP file") # Get the top-level label, which should represent an assembly top_level_label = labels.Value(1) - # Make sure there is a top-level assembly if shape_tool.IsTopLevel(top_level_label) and shape_tool.IsAssembly_s( top_level_label @@ -221,19 +198,16 @@ def _process_label(lbl: TDF_Label): # Start the recursive processing of labels imported_assy = _process_label(top_level_label) - if imported_assy and hasattr(imported_assy, "children"): - # Check to see if there is an extra top-level node. This is done because - # cq.Assembly.export adds an extra top-level node which will cause a cascade of - # extras on successive round-trips. exportStepMeta does not add the extra top-level - # node and so does not exhibit this behavior. - if assy.name == imported_assy.children[0].name: - imported_assy = imported_assy.children[0] - - # Copy all of the children over to the main assembly object - for child in imported_assy.children: - assy.add(child, name=child.name, color=child.color, loc=child.loc) - elif imported_assy: - assy.add(imported_assy) + # Handle a possible extra top-level node. This is done because cq.Assembly.export + # adds an extra top-level node which will cause a cascade of + # extras on successive round-trips. exportStepMeta does not add the extra top-level + # node and so does not exhibit this behavior. + if assy.name == imported_assy.children[0].name: + imported_assy = imported_assy.children[0] + + # Copy all of the children over to the main assembly object + for child in imported_assy.children: + assy.add(child, name=child.name, color=child.color, loc=child.loc) # Copy across subshape data for shape, name in imported_assy._subshape_names.items(): From 7da1eab99c93b3a2f05a4d372da4d8154e1e9b42 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Wed, 30 Jul 2025 15:26:30 -0400 Subject: [PATCH 42/44] More cleanup to try to get code coverage high enough without creating a contrived STEP file --- cadquery/occ_impl/importers/assembly.py | 223 +++++++++++------------- 1 file changed, 106 insertions(+), 117 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 62b4b212c..431d2b4a8 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -5,9 +5,9 @@ from OCP.IFSelect import IFSelect_RetDone from OCP.TDocStd import TDocStd_Document from OCP.TDataStd import TDataStd_Name +from OCP.TNaming import TNaming_NamedShape from OCP.STEPCAFControl import STEPCAFControl_Reader -from OCP.XCAFDoc import XCAFDoc_ColorSurf -from OCP.XCAFDoc import XCAFDoc_DocumentTool +from OCP.XCAFDoc import XCAFDoc_ColorSurf, XCAFDoc_DocumentTool, XCAFDoc_GraphNode import cadquery as cq from ..assembly import AssemblyProtocol @@ -27,135 +27,124 @@ def _process_label(lbl: TDF_Label): """ Recursive method to process the assembly in a top-down manner. """ - # If we have an assembly, extract all of the information out of it that we can - if shape_tool.IsAssembly_s(lbl): - # Instantiate the new assembly - new_assy = cq.Assembly() - - # Look for components - comp_labels = TDF_LabelSequence() - shape_tool.GetComponents_s(lbl, comp_labels) - - for i in range(comp_labels.Length()): - comp_label = comp_labels.Value(i + 1) - - # Get the location of the component label - loc = shape_tool.GetLocation_s(comp_label) - cq_loc = cq.Location(loc) if loc else None - - if shape_tool.IsReference_s(comp_label): - ref_label = TDF_Label() - shape_tool.GetReferredShape_s(comp_label, ref_label) - - # Find the name of this referenced part - ref_name_attr = TDataStd_Name() - if ref_label.FindAttribute(TDataStd_Name.GetID_s(), ref_name_attr): - ref_name = str(ref_name_attr.Get().ToExtString()) - - if shape_tool.IsAssembly_s(ref_label): - # Recursively process subassemblies - sub_assy = _process_label(ref_label) - - # Add the appropriate attributes to the subassembly - new_assy.add(sub_assy, name=f"{ref_name}", loc=cq_loc) - elif shape_tool.IsSimpleShape_s(ref_label): - # A single shape needs to be added to the assembly - final_shape = shape_tool.GetShape_s(ref_label) - cq_shape = cq.Shape.cast(final_shape) + + # Instantiate the new assembly + new_assy = cq.Assembly() + + # Look for components + comp_labels = TDF_LabelSequence() + shape_tool.GetComponents_s(lbl, comp_labels) + + for i in range(comp_labels.Length()): + comp_label = comp_labels.Value(i + 1) + + # Get the location of the component label + loc = shape_tool.GetLocation_s(comp_label) + cq_loc = cq.Location(loc) if loc else None + + if shape_tool.IsReference_s(comp_label): + ref_label = TDF_Label() + shape_tool.GetReferredShape_s(comp_label, ref_label) + + # Find the name of this referenced part + ref_name_attr = TDataStd_Name() + if ref_label.FindAttribute(TDataStd_Name.GetID_s(), ref_name_attr): + ref_name = str(ref_name_attr.Get().ToExtString()) + + if shape_tool.IsAssembly_s(ref_label): + # Recursively process subassemblies + sub_assy = _process_label(ref_label) + + # Add the appropriate attributes to the subassembly + new_assy.add(sub_assy, name=f"{ref_name}", loc=cq_loc) + elif shape_tool.IsSimpleShape_s(ref_label): + # A single shape needs to be added to the assembly + final_shape = shape_tool.GetShape_s(ref_label) + cq_shape = cq.Shape.cast(final_shape) + + # Find the subshape color, if there is one set for this shape + color = Quantity_ColorRGBA() + # Extract the color, if present on the shape + if color_tool.GetColor(final_shape, XCAFDoc_ColorSurf, color): + rgb = color.GetRGB() + cq_color = cq.Color( + rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha() + ) + else: + cq_color = None + + new_assy.add( + cq_shape, name=f"{ref_name}", loc=cq_loc, color=cq_color + ) + + # Search for subshape names, layers and colors + for j in range(ref_label.NbChildren()): + child_label = ref_label.FindChild(j + 1) + + # Save the shape so that we can add it to the subshape data + cur_shape: TopoDS_Shape = shape_tool.GetShape_s(child_label) + + # Find the layer name, if there is one set for this shape + layers = TDF_LabelSequence() + layer_tool.GetLayers(child_label, layers) + for i in range(1, layers.Length() + 1): + lbl = layers.Value(i) + name_attr = TDataStd_Name() + lbl.FindAttribute(TDataStd_Name.GetID_s(), name_attr) + + # Extract the layer name for the shape here + layer_name = name_attr.Get().ToExtString() + + # Add the layer as a subshape entry on the assembly + new_assy.addSubshape( + cq.Shape.cast(cur_shape), layer=layer_name + ) # Find the subshape color, if there is one set for this shape color = Quantity_ColorRGBA() # Extract the color, if present on the shape - if color_tool.GetColor(final_shape, XCAFDoc_ColorSurf, color): + if color_tool.GetColor(cur_shape, XCAFDoc_ColorSurf, color): rgb = color.GetRGB() cq_color = cq.Color( - rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha() + rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha(), ) - else: - cq_color = None - - new_assy.add( - cq_shape, name=f"{ref_name}", loc=cq_loc, color=cq_color - ) - - # Search for subshape names, layers and colors - for j in range(ref_label.NbChildren()): - child_label = ref_label.FindChild(j + 1) - # Save the shape so that we can add it to the subshape data - cur_shape: TopoDS_Shape = shape_tool.GetShape_s(child_label) + # Save the color info via the assembly subshape mechanism + new_assy.addSubshape( + cq.Shape.cast(cur_shape), color=cq_color + ) - # Find the layer name, if there is one set for this shape - layers = TDF_LabelSequence() - layer_tool.GetLayers(child_label, layers) - for i in range(1, layers.Length() + 1): - lbl = layers.Value(i) + # Iterate through all the attributes looking for subshape names. + # This is safer than trying to access the attributes directly with + # FindAttribute because it will cause a segfault in certain cases. + attr_iterator = TDF_AttributeIterator(child_label) + while attr_iterator.More(): + current_attr = attr_iterator.Value() + + # TNaming_NamedShape is used to store and manage references to + # topological shapes, and its attributes can be accessed directly. + # XCAFDoc_GraphNode contains a graph of labels, and so we must + # follow the branch back to a father. + if isinstance(current_attr, XCAFDoc_GraphNode): + lbl = current_attr.GetFather(1).Label() + + # Find the name attribute and add it for the subshape name_attr = TDataStd_Name() - lbl.FindAttribute(TDataStd_Name.GetID_s(), name_attr) - - # Extract the layer name for the shape here - layer_name = name_attr.Get().ToExtString() - - # Add the layer as a subshape entry on the assembly - new_assy.addSubshape( - cq.Shape.cast(cur_shape), layer=layer_name - ) - - # Find the subshape color, if there is one set for this shape - color = Quantity_ColorRGBA() - # Extract the color, if present on the shape - if color_tool.GetColor(cur_shape, XCAFDoc_ColorSurf, color): - rgb = color.GetRGB() - cq_color = cq.Color( - rgb.Red(), rgb.Green(), rgb.Blue(), color.Alpha(), - ) - - # Save the color info via the assembly subshape mechanism - new_assy.addSubshape( - cq.Shape.cast(cur_shape), color=cq_color - ) - - # Iterate through all the attributes looking for subshape names. - # This is safer than trying to access the attributes directly with - # FindAttribute because it will cause a segfault in certain cases. - attr_iterator = TDF_AttributeIterator(child_label) - while attr_iterator.More(): - current_attr = attr_iterator.Value() - - # TNaming_NamedShape is used to store and manage references to - # topological shapes, and its attributes can be accessed directly. - # XCAFDoc_GraphNode contains a graph of labels, and so we must - # follow the branch back to a father. - if ( - current_attr.DynamicType().Name() - == "XCAFDoc_GraphNode" - ): - # Only the GraphNode should have this method - if hasattr(current_attr, "GetFather"): - lbl = current_attr.GetFather(1).Label() - - # Find the name attribute and add it for the subshape - name_attr = TDataStd_Name() - if lbl.FindAttribute( - TDataStd_Name.GetID_s(), name_attr - ): - # Save this as the name of the subshape - new_assy.addSubshape( - cq.Shape.cast(cur_shape), - name=name_attr.Get().ToExtString(), - ) - elif ( - current_attr.DynamicType().Name() - == "TNaming_NamedShape" + if lbl.FindAttribute( + TDataStd_Name.GetID_s(), name_attr ): - # Save the shape so that we can add it to the subshape data - cur_shape: TopoDS_Shape = shape_tool.GetShape_s( - child_label + # Save this as the name of the subshape + new_assy.addSubshape( + cq.Shape.cast(cur_shape), + name=name_attr.Get().ToExtString(), ) + elif isinstance(current_attr, TNaming_NamedShape): + # Save the shape so that we can add it to the subshape data + cur_shape = shape_tool.GetShape_s(child_label) - attr_iterator.Next() + attr_iterator.Next() - return new_assy + return new_assy # Document that the step file will be read into doc = TDocStd_Document(TCollection_ExtendedString("XmlOcaf")) From 8251d92be2d5a26e2b85afa5c3bf58d1612b4a8b Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Mon, 4 Aug 2025 11:32:36 -0400 Subject: [PATCH 43/44] Fix lack of application to the top-level assembly --- cadquery/occ_impl/importers/assembly.py | 8 ++++++++ tests/test_assembly.py | 5 ++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index 431d2b4a8..ffc1023d3 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -184,6 +184,14 @@ def _process_label(lbl: TDF_Label): top_level_label.FindAttribute(TDataStd_Name.GetID_s(), name_attr) assy.name = str(name_attr.Get().ToExtString()) + # Get the location of the top-level component + comp_labels = TDF_LabelSequence() + shape_tool.GetComponents_s(top_level_label, comp_labels) + comp_label = comp_labels.Value(1) + loc = shape_tool.GetLocation_s(comp_label) + if loc and not loc.IsIdentity(): + assy.loc = cq.Location(loc) + # Start the recursive processing of labels imported_assy = _process_label(top_level_label) diff --git a/tests/test_assembly.py b/tests/test_assembly.py index f045671d5..c33b6e361 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -996,7 +996,7 @@ def test_plain_assembly_import(tmp_path_factory): cube_3 = cq.Workplane().box(5, 5, 5) cube_4 = cq.Workplane().box(5, 5, 5) - assy = cq.Assembly(name="top_level") + assy = cq.Assembly(name="top_level", loc=cq.Location(10, 10, 10)) assy.add(cube_1, color=cq.Color("green")) assy.add(cube_2, loc=cq.Location((10, 10, 10)), color=cq.Color("red")) assy.add(cube_3, loc=cq.Location((-10, -10, -10)), color=cq.Color("red")) @@ -1015,6 +1015,9 @@ def test_plain_assembly_import(tmp_path_factory): assert imported_assy.children[2].loc.toTuple()[0] == (-10.0, -10.0, -10.0,) assert imported_assy.children[3].loc.toTuple()[0] == (10.0, -10.0, -10.0,) + # Make sure the location of the top-level assembly was preserved + assert imported_assy.loc.toTuple() == cq.Location((10, 10, 10)).toTuple() + # Check the colors assert pytest.approx( imported_assy.children[0].children[0].color.toTuple(), rel=0.01 From 1543908cc516039a85752137a2c7fe80731df474 Mon Sep 17 00:00:00 2001 From: Jeremy Wright Date: Mon, 4 Aug 2025 13:07:14 -0400 Subject: [PATCH 44/44] Trying to increase test coverage high enough --- cadquery/occ_impl/importers/assembly.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cadquery/occ_impl/importers/assembly.py b/cadquery/occ_impl/importers/assembly.py index ffc1023d3..182578160 100644 --- a/cadquery/occ_impl/importers/assembly.py +++ b/cadquery/occ_impl/importers/assembly.py @@ -189,8 +189,7 @@ def _process_label(lbl: TDF_Label): shape_tool.GetComponents_s(top_level_label, comp_labels) comp_label = comp_labels.Value(1) loc = shape_tool.GetLocation_s(comp_label) - if loc and not loc.IsIdentity(): - assy.loc = cq.Location(loc) + assy.loc = cq.Location(loc) # Start the recursive processing of labels imported_assy = _process_label(top_level_label)