STEP Export with Face Metadata (Names and Colors) (#1782)

jmwright · adam-urbanczyk · web-flow · commit db62a1d71a4a · 2025-04-16T08:14:33.000-04:00
* Moved name, color and layer metadata into cq.Assembly

* Started trying to exert more control over the internal structure of the STEP structure

* Working implementation that works with deeply nested assemblies

* Added write_pcurves setting to increase test coverage

* Added mention of the metadata export method in the documentation

* Removed extra properties from docstring

---------

Co-authored-by: AU &lt;adam-urbanczyk@users.noreply.github.com&gt;
diff --git a/cadquery/assembly.py b/cadquery/assembly.py
@@ -94,6 +94,11 @@ class Assembly(object):
     objects: Dict[str, "Assembly"]
     constraints: List[Constraint]
 
+    # Allows metadata to be stored for exports
+    _subshape_names: dict[Shape, str]
+    _subshape_colors: dict[Shape, Color]
+    _subshape_layers: dict[Shape, str]
+
     _solve_result: Optional[Dict[str, Any]]
 
     def __init__(
@@ -139,6 +144,10 @@ def __init__(
 
         self._solve_result = None
 
+        self._subshape_names = {}
+        self._subshape_colors = {}
+        self._subshape_layers = {}
+
     def _copy(self) -> "Assembly":
         """
         Make a deep copy of an assembly
@@ -685,3 +694,30 @@ def _repr_javascript_(self):
         from .occ_impl.jupyter_tools import display
 
         return display(self)._repr_javascript_()
+
+    def addSubshape(
+        self,
+        s: Shape,
+        name: Optional[str] = None,
+        color: Optional[Color] = None,
+        layer: Optional[str] = None,
+    ) -> "Assembly":
+        """
+        Handles name, color and layer metadata for subshapes.
+
+        :param s: The subshape to add metadata to.
+        :param name: The name to assign to the subshape.
+        :param color: The color to assign to the subshape.
+        :param layer: The layer to assign to the subshape.
+        :return: The modified assembly.
+        """
+
+        # Handle any metadata we were passed
+        if name:
+            self._subshape_names[s] = name
+        if color:
+            self._subshape_colors[s] = color
+        if layer:
+            self._subshape_layers[s] = layer
+
+        return self
diff --git a/cadquery/occ_impl/assembly.py b/cadquery/occ_impl/assembly.py
@@ -160,6 +160,18 @@ def shapes(self) -> Iterable[Shape]:
     def children(self) -> Iterable["AssemblyProtocol"]:
         ...
 
+    @property
+    def _subshape_names(self) -> Dict[Shape, str]:
+        ...
+
+    @property
+    def _subshape_colors(self) -> Dict[Shape, Color]:
+        ...
+
+    @property
+    def _subshape_layers(self) -> Dict[Shape, str]:
+        ...
+
     def traverse(self) -> Iterable[Tuple[str, "AssemblyProtocol"]]:
         ...
 
diff --git a/cadquery/occ_impl/exporters/assembly.py b/cadquery/occ_impl/exporters/assembly.py
@@ -14,7 +14,11 @@
 from OCP.STEPCAFControl import STEPCAFControl_Writer
 from OCP.STEPControl import STEPControl_StepModelType
 from OCP.IFSelect import IFSelect_ReturnStatus
+from OCP.TDF import TDF_Label
+from OCP.TDataStd import TDataStd_Name
+from OCP.TDocStd import TDocStd_Document
 from OCP.XCAFApp import XCAFApp_Application
+from OCP.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ColorGen
 from OCP.XmlDrivers import (
     XmlDrivers_DocumentStorageDriver,
     XmlDrivers_DocumentRetrievalDriver,
@@ -28,6 +32,8 @@
 
 from ..assembly import AssemblyProtocol, toCAF, toVTK, toFusedCAF
 from ..geom import Location
+from ..shapes import Shape, Compound
+from ..assembly import Color
 
 
 class ExportModes:
@@ -42,7 +48,7 @@ def exportAssembly(
     assy: AssemblyProtocol,
     path: str,
     mode: STEPExportModeLiterals = "default",
-    **kwargs
+    **kwargs,
 ) -> bool:
     """
     Export an assembly to a STEP file.
@@ -99,6 +105,168 @@ def exportAssembly(
     return status == IFSelect_ReturnStatus.IFSelect_RetDone
 
 
+def exportStepMeta(
+    assy: AssemblyProtocol,
+    path: str,
+    write_pcurves: bool = True,
+    precision_mode: int = 0,
+) -> bool:
+    """
+    Export an assembly to a STEP file with faces tagged with names and colors. This is done as a
+    separate method from the main STEP export because this is not compatible with the fused mode
+    and also flattens the hierarchy of the STEP.
+
+    Layers are used because some software does not understand the ADVANCED_FACE entity and needs
+    names attached to layers instead.
+
+    :param assy: assembly
+    :param path: Path and filename for writing
+    :param write_pcurves: Enable or disable writing parametric curves to the STEP file. Default True.
+        If False, writes STEP file without pcurves. This decreases the size of the resulting STEP file.
+    :param precision_mode: Controls the uncertainty value for STEP entities. Specify -1, 0, or 1. Default 0.
+        See OCCT documentation.
+    """
+
+    pcurves = 1
+    if not write_pcurves:
+        pcurves = 0
+
+    # Initialize the XCAF document that will allow the STEP export
+    app = XCAFApp_Application.GetApplication_s()
+    doc = TDocStd_Document(TCollection_ExtendedString("XmlOcaf"))
+    app.InitDocument(doc)
+
+    # Shape and color tools
+    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_child(child: AssemblyProtocol, assy_label: TDF_Label):
+        """
+        Process a child part which is not a subassembly.
+        :param child: Child part to process (we should already have filtered out subassemblies)
+        :param assy_label: The label for the assembly to add this part to
+        :return: None
+        """
+
+        child_items = None
+
+        # We combine these because the metadata could be stored at the parent or child level
+        combined_names = {**assy._subshape_names, **child._subshape_names}
+        combined_colors = {**assy._subshape_colors, **child._subshape_colors}
+        combined_layers = {**assy._subshape_layers, **child._subshape_layers}
+
+        # Collect all of the shapes in the child object
+        if child.obj:
+            child_items = (
+                child.obj
+                if isinstance(child.obj, Shape)
+                else Compound.makeCompound(
+                    s for s in child.obj.vals() if isinstance(s, Shape)
+                ),
+                child.name,
+                child.loc,
+                child.color,
+            )
+
+        if child_items:
+            shape, name, loc, color = child_items
+
+            # Handle shape name, color and location
+            part_label = shape_tool.AddShape(shape.wrapped, False)
+            TDataStd_Name.Set_s(part_label, TCollection_ExtendedString(name))
+            if color:
+                color_tool.SetColor(part_label, color.wrapped, XCAFDoc_ColorGen)
+            shape_tool.AddComponent(assy_label, part_label, loc.wrapped)
+
+            # If this assembly has shape metadata, add it to the shape
+            if (
+                len(combined_names) > 0
+                or len(combined_colors) > 0
+                or len(combined_layers) > 0
+            ):
+                names = combined_names
+                colors = combined_colors
+                layers = combined_layers
+
+                # Step through every face in the shape, and see if any metadata needs to be attached to it
+                for face in shape.Faces():
+                    if face in names or face in shape in colors or face in layers:
+                        # Add the face as a subshape
+                        face_label = shape_tool.AddSubShape(part_label, face.wrapped)
+
+                        # In some cases the face may not be considered part of the shape, so protect
+                        # against that
+                        if not face_label.IsNull():
+                            # Set the ADVANCED_FACE label, even though the layer holds the same data
+                            if face in names:
+                                TDataStd_Name.Set_s(
+                                    face_label, TCollection_ExtendedString(names[face])
+                                )
+
+                            # Set the individual face color
+                            if face in colors:
+                                color_tool.SetColor(
+                                    face_label, colors[face].wrapped, XCAFDoc_ColorGen,
+                                )
+
+                            # Also add a layer to hold the face label data
+                            if face in layers:
+                                layer_label = layer_tool.AddLayer(
+                                    TCollection_ExtendedString(layers[face])
+                                )
+                                layer_tool.SetLayer(face_label, layer_label)
+
+    def _process_assembly(
+        assy: AssemblyProtocol, parent_label: Optional[TDF_Label] = None
+    ):
+        """
+        Recursively process the assembly and its children.
+        :param assy: Assembly to process
+        :param parent_label: The parent label for the assembly
+        :return: None
+        """
+        # Use the assembly name if the user set it
+        assembly_name = assy.name if assy.name else str(uuid.uuid1())
+
+        # Create the top level object that will hold all the subassemblies and parts
+        assy_label = shape_tool.NewShape()
+        TDataStd_Name.Set_s(assy_label, TCollection_ExtendedString(assembly_name))
+
+        # Handle subassemblies
+        if parent_label:
+            shape_tool.AddComponent(parent_label, assy_label, assy.loc.wrapped)
+
+        # The children may be parts or assemblies
+        for child in assy.children:
+            # Child is a part
+            if len(list(child.children)) == 0:
+                _process_child(child, assy_label)
+            # Child is a subassembly
+            else:
+                _process_assembly(child, assy_label)
+
+    _process_assembly(assy)
+
+    # Update the assemblies
+    shape_tool.UpdateAssemblies()
+
+    # Set up the writer and write the STEP file
+    session = XSControl_WorkSession()
+    writer = STEPCAFControl_Writer(session, False)
+    Interface_Static.SetIVal_s("write.stepcaf.subshapes.name", 1)
+    writer.SetColorMode(True)
+    writer.SetLayerMode(True)
+    writer.SetNameMode(True)
+    Interface_Static.SetIVal_s("write.surfacecurve.mode", pcurves)
+    Interface_Static.SetIVal_s("write.precision.mode", precision_mode)
+    writer.Transfer(doc, STEPControl_StepModelType.STEPControl_AsIs)
+
+    status = writer.Write(path)
+
+    return status == IFSelect_ReturnStatus.IFSelect_RetDone
+
+
 def exportCAF(assy: AssemblyProtocol, path: str) -> bool:
     """
     Export an assembly to a OCAF xml file (internal OCCT format).
diff --git a/doc/importexport.rst b/doc/importexport.rst
@@ -205,6 +205,32 @@ This is done by setting the name property of the assembly before calling :meth:`
 
 If an assembly name is not specified, a UUID will be used to avoid name conflicts.
 
+Exporting Assemblies to STEP with Metadata
+###########################################
+
+It is possible to attach metadata to the assembly that will be included in the STEP file. This metadata can be attached to arbitrary shapes and includes names, colors and layers. This is done by using the :meth:`Assembly.addSubshape` method before calling :meth:`cadquery.occ_impl.exporters.assembly.exportStepMeta`.
+
+.. code-block:: python
+
+    import cadquery as cq
+    from cadquery.occ_impl.exporters.assembly import exportStepMeta
+
+    # 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"
+    )
+
+    # Export the assembly to STEP with metadata
+    exportStepMeta(assy, "out.step")
+
 Exporting Assemblies to glTF
 #############################
 
diff --git a/tests/test_assembly.py b/tests/test_assembly.py
