Materials Support for STEP Import/Export (#1956)

* Added material import for round-trip STEP export/import

* Added None to _get_material call

* Parameterized the test for other formats

* Fixed default material description and comment

---------

Co-authored-by: adam-urbanczyk <13981538+adam-urbanczyk@users.noreply.github.com>

Author: jmwright

cadquery/assembly.py

@@ -195,6 +195,7 @@ def add(
         loc: Optional[Location] = None,
         name: Optional[str] = None,
         color: Optional[Color] = None,
+        material: Optional[Union[Material, str]] = None,
     ) -> Self:
         """
         Add a subassembly to the current assembly.

cadquery/occ_impl/assembly.py

@@ -83,7 +83,7 @@ def __init__(self, name: str | None = None, **kwargs):
 
         # Default values in case the user did not set any others
         aName = "Default"
-        aDescription = "Default material with properties similar to low carbon steel"
+        aDescription = ""
         aDensity = 7.85
         aDensityName = "Mass density"
         aDensityTypeName = "g/cm^3"
@@ -285,6 +285,7 @@ def __init__(
         loc: Optional[Location] = None,
         name: Optional[str] = None,
         color: Optional[Color] = None,
+        material: Optional[Material] = None,
     ):
         ...
 
@@ -316,6 +317,14 @@ def color(self) -> Optional[Color]:
     def color(self, value: Optional[Color]) -> None:
         ...
 
+    @property
+    def material(self) -> Optional[Material]:
+        ...
+
+    @material.setter
+    def material(self, value: Optional[Material]) -> None:
+        ...
+
     @property
     def obj(self) -> AssemblyObjects:
         ...
@@ -355,6 +364,7 @@ def add(
         loc: Optional[Location] = None,
         name: Optional[str] = None,
         color: Optional[Color] = None,
+        material: Optional[Union[Material, str]] = None,
     ) -> Self:
         ...
 
@@ -422,6 +432,18 @@ def setColor(l: TDF_Label, color: Color, tool):
     tool.SetColor(l, color.wrapped, XCAFDoc_ColorType.XCAFDoc_ColorSurf)
 
 
+def setMaterial(l: TDF_Label, material: Material, tool):
+
+    tool.SetMaterial(
+        l,
+        TCollection_HAsciiString(material.name),
+        TCollection_HAsciiString(material.description),
+        material.density,
+        TCollection_HAsciiString("MassDensity"),
+        TCollection_HAsciiString(material.densityUnit),
+    )
+
+
 def toCAF(
     assy: AssemblyProtocol,
     coloredSTEP: bool = False,
@@ -447,6 +469,7 @@ def toCAF(
     tool.SetAutoNaming_s(False)
     ctool = XCAFDoc_DocumentTool.ColorTool_s(doc.Main())
     ltool = XCAFDoc_DocumentTool.LayerTool_s(doc.Main())
+    mtool = XCAFDoc_DocumentTool.MaterialTool_s(doc.Main())
 
     # used to store labels with unique part-color combinations
     unique_objs: Dict[Tuple[Color | None, AssemblyObjects], TDF_Label] = {}
@@ -463,6 +486,9 @@ def _toCAF(el: AssemblyProtocol, ancestor: TDF_Label | None) -> TDF_Label:
         # define the current color
         current_color = el.color if el.color else None
 
+        # define the current material
+        current_material = el.material if el.material else None
+
         # add a leaf with the actual part if needed
         if el.obj:
             # get/register unique parts referenced in the assy
@@ -490,6 +516,10 @@ def _toCAF(el: AssemblyProtocol, ancestor: TDF_Label | None) -> TDF_Label:
                 if coloredSTEP and current_color:
                     setColor(lab, current_color, ctool)
 
+                # Handle materials when exporting to STEP
+                if current_material:
+                    setMaterial(lab, current_material, mtool)
+
             # handle subshape names/colors/layers
             subshape_colors = el._subshape_colors
             subshape_names = el._subshape_names

cadquery/occ_impl/importers/assembly.py

@@ -15,14 +15,15 @@
     XCAFDoc_ColorTool,
     XCAFDoc,
     XCAFDoc_ColorType,
+    XCAFDoc_Material,
 )
 from OCP.TDocStd import TDocStd_Application
 from OCP.XmlXCAFDrivers import XmlXCAFDrivers
 from OCP.BinXCAFDrivers import BinXCAFDrivers
 from OCP.Interface import Interface_Static
 from OCP.PCDM import PCDM_ReaderStatus
 
-from ..assembly import AssemblyProtocol, Color
+from ..assembly import AssemblyProtocol, Color, Material
 from ..geom import Location
 from ..shapes import Shape
 
@@ -42,6 +43,38 @@ def _get_name(label: TDF_Label) -> str:
     return rv
 
 
+def _get_material(label: TDF_Label) -> Material | None:
+    """
+    Helper to get the material for a given label.
+    """
+
+    rv = None
+
+    material_ref_guid = XCAFDoc.MaterialRefGUID_s()
+
+    if label.IsAttribute(material_ref_guid):
+        attr = TDataStd_TreeNode()
+        label.FindAttribute(material_ref_guid, attr)
+        material_label = attr.Father().Label()
+
+        material_attr = XCAFDoc_Material()
+
+        material_label.FindAttribute(XCAFDoc_Material.GetID_s(), material_attr)
+        name = material_attr.GetName().ToCString()
+        description = material_attr.GetDescription().ToCString()
+        density = material_attr.GetDensity()
+        density_unit = material_attr.GetDensValType().ToCString()
+
+        rv = Material(
+            name=name,
+            description=description,
+            density=density,
+            densityUnit=density_unit,
+        )
+
+    return rv
+
+
 def _get_ref_color(label: TDF_Label) -> Color | None:
     """
     Helper to get the instance color of a given label.
@@ -211,6 +244,7 @@ def _process_label(lbl: TDF_Label, parent: AssemblyProtocol):
 
                 # get (if it exists the color of the comp label)
                 color = _get_ref_color(comp_label)
+                material = _get_material(comp_label)
 
                 if shape_tool.IsAssembly_s(ref_label):
                     # Find the name of this referenced part
@@ -222,7 +256,13 @@ def _process_label(lbl: TDF_Label, parent: AssemblyProtocol):
                     _ = _process_label(ref_label, sub_assy)
 
                     # Add the subassy
-                    parent.add(sub_assy, name=ref_name, loc=cq_loc, color=color)
+                    parent.add(
+                        sub_assy,
+                        loc=cq_loc,
+                        name=ref_name,
+                        color=color,
+                        material=material,
+                    )
 
                 elif shape_tool.IsSimpleShape_s(ref_label):
                     # Find the name of this referenced part
@@ -236,6 +276,9 @@ def _process_label(lbl: TDF_Label, parent: AssemblyProtocol):
                     if color is None:
                         color = _get_shape_color(final_shape, color_tool)
 
+                    if material is None:
+                        material = _get_material(ref_label)
+
                     # this if/else is needed to handle different structures of STEP files
                     # "*"/"*_part" based naming is the default structure produced by CQ
                     # with an object and child nodes at the same time
@@ -248,7 +291,11 @@ def _process_label(lbl: TDF_Label, parent: AssemblyProtocol):
                         current = parent
                     else:
                         tmp = assy.__class__(
-                            cq_shape, loc=cq_loc, name=comp_name, color=color
+                            cq_shape,
+                            loc=cq_loc,
+                            name=comp_name,
+                            color=color,
+                            material=material,
                         )
                         parent.add(tmp)
 
@@ -287,6 +334,7 @@ def _process_label(lbl: TDF_Label, parent: AssemblyProtocol):
 
                         # try the instance first
                         color = _get_ref_color(child_label)
+                        material = _get_material(child_label)
 
                         if color:
                             # Save the color info via the assembly subshape mechanism

tests/test_assembly.py

@@ -878,6 +878,41 @@ def test_meta_step_export_edge_cases(tmp_path_factory):
     assert success
 
 
+@pytest.mark.parametrize("kind", ["step", "xml", "xbf"])
+def test_step_roundtrip_with_materials(kind, tmp_path_factory):
+    """
+    Tests to make sure that once materials have been exported to a file format
+    such as STEP, the materials can be imported again.
+    """
+
+    materials_assy = cq.Assembly()
+    materials_assy.add(
+        cq.Workplane().box(10, 10, 10),
+        name="cube_1",
+        material=cq.Material(name="copper"),
+    )
+
+    # Use a temporary directory
+    tmpdir = tmp_path_factory.mktemp("out")
+    materials_path = os.path.join(tmpdir, f"roundtrip_materials.{kind}")
+
+    exportAssembly(materials_assy, materials_path)
+
+    # Read the contents as a step file as a string so we can check the outputs
+    with open(materials_path, "r") as f:
+        step_contents = f.read()
+
+        # Make sure that the face name string is present in the exported STEP contents
+        assert "copper" in step_contents
+
+    # Import the STEP file back in as an assembly
+    test_assy = cq.Assembly.importStep(materials_path)
+
+    # Make sure that the material was re-imported
+    assert test_assy.children[0].material is not None
+    assert test_assy.children[0].material.name == "copper"
+
+
 def test_assembly_step_import(tmp_path_factory, subshape_assy):
     """
     Test if the STEP import works correctly for an assembly with subshape data attached.
@@ -1545,10 +1580,7 @@ def test_materials():
     mat_1 = cq.Material()
     assy.add(wp_1, material=mat_1)
     assert assy.children[0].material.name == "Default"
-    assert (
-        assy.children[0].material.description
-        == "Default material with properties similar to low carbon steel"
-    )
+    assert assy.children[0].material.description == ""
     assert assy.children[0].material.density == 7.85
     assert assy.children[0].material.densityUnit == "g/cm^3"
 
@@ -1568,7 +1600,7 @@ def test_materials():
     # Test the ability to convert a material to a tuple
     assert mat_2.toTuple() == ("test", "Test material", 1.0, "lb/in^3")
 
-    # Test the ability to has a material
+    # Test the ability to have a material
     assert mat_2.__hash__() == hash(("test", "Test material", 1.0, "lb/in^3"))
 
     # Test the equality operator with material