Add Basic Materials Functionality (#1923)

jmwright · web-flow · commit 2e0581991f98 · 2025-12-12T12:46:08.000-05:00
* Allow passing of material as a string when adding a subassembly

* Fixed metadata bug because of _copy call

* Reworked Material class to wrap XCAFDoc_Material and XCAFDoc_VisMaterial instead

* Allow pickling of material class
diff --git a/cadquery/__init__.py b/cadquery/__init__.py
@@ -37,7 +37,7 @@
 )
 from .sketch import Sketch
 from .cq import CQ, Workplane
-from .assembly import Assembly, Color, Constraint
+from .assembly import Assembly, Color, Constraint, Material
 from . import selectors
 from . import plugins
 
@@ -48,6 +48,7 @@
     "Assembly",
     "Color",
     "Constraint",
+    "Material",
     "plugins",
     "selectors",
     "Plane",
diff --git a/cadquery/assembly.py b/cadquery/assembly.py
@@ -20,7 +20,7 @@
 from .cq import Workplane
 from .occ_impl.shapes import Shape, Compound, isSubshape, compound
 from .occ_impl.geom import Location
-from .occ_impl.assembly import Color
+from .occ_impl.assembly import Color, Material
 from .occ_impl.solver import (
     ConstraintKind,
     ConstraintSolver,
@@ -85,12 +85,20 @@ def _define_grammar():
 _grammar = _define_grammar()
 
 
+def _ensure_material(material):
+    """
+    Convert string to Material if needed.
+    """
+    return Material(material) if isinstance(material, str) else material
+
+
 class Assembly(object):
     """Nested assembly of Workplane and Shape objects defining their relative positions."""
 
     loc: Location
     name: str
     color: Optional[Color]
+    material: Optional[Material]
     metadata: Dict[str, Any]
 
     obj: AssemblyObjects
@@ -113,6 +121,7 @@ def __init__(
         loc: Optional[Location] = None,
         name: Optional[str] = None,
         color: Optional[Color] = None,
+        material: Optional[Material] = None,
         metadata: Optional[Dict[str, Any]] = None,
     ):
         """
@@ -122,6 +131,7 @@ def __init__(
         :param loc: location of the root object (default: None, interpreted as identity transformation)
         :param name: unique name of the root object (default: None, resulting in an UUID being generated)
         :param color: color of the added object (default: None)
+        :param material: material (for visual and/or physical properties) of the added object (default: None)
         :param metadata: a store for user-defined metadata (default: None)
         :return: An Assembly object.
 
@@ -141,6 +151,7 @@ def __init__(
         self.loc = loc if loc else Location()
         self.name = name if name else str(uuid())
         self.color = color if color else None
+        self.material = material if material else None
         self.metadata = metadata if metadata else {}
         self.parent = None
 
@@ -159,7 +170,9 @@ def _copy(self) -> "Assembly":
         Make a deep copy of an assembly
         """
 
-        rv = self.__class__(self.obj, self.loc, self.name, self.color, self.metadata)
+        rv = self.__class__(
+            self.obj, self.loc, self.name, self.color, self.material, self.metadata
+        )
 
         rv._subshape_colors = BiDict(self._subshape_colors)
         rv._subshape_names = BiDict(self._subshape_names)
@@ -203,6 +216,7 @@ def add(
         loc: Optional[Location] = None,
         name: Optional[str] = None,
         color: Optional[Color] = None,
+        material: Optional[Union[Material, str]] = None,
         metadata: Optional[Dict[str, Any]] = None,
     ) -> Self:
         """
@@ -214,6 +228,8 @@ def add(
         :param name: unique name of the root object (default: None, resulting in an UUID being
           generated)
         :param color: color of the added object (default: None)
+        :param material: material (for visual and/or physical properties) of the added object
+          (default: None)
         :param metadata: a store for user-defined metadata (default: None)
         """
         ...
@@ -237,15 +253,23 @@ def add(self, arg, **kwargs):
             subassy.loc = kwargs["loc"] if kwargs.get("loc") else arg.loc
             subassy.name = kwargs["name"] if kwargs.get("name") else arg.name
             subassy.color = kwargs["color"] if kwargs.get("color") else arg.color
+            subassy.material = _ensure_material(
+                kwargs["material"] if kwargs.get("material") else arg.material
+            )
             subassy.metadata = (
                 kwargs["metadata"] if kwargs.get("metadata") else arg.metadata
             )
+
             subassy.parent = self
 
             self.children.append(subassy)
             self.objects.update(subassy._flatten())
 
         else:
+            # Convert the material string to a Material object, if needed
+            if "material" in kwargs:
+                kwargs["material"] = _ensure_material(kwargs["material"])
+
             assy = self.__class__(arg, **kwargs)
             assy.parent = self
 
diff --git a/cadquery/occ_impl/assembly.py b/cadquery/occ_impl/assembly.py
@@ -13,12 +13,15 @@
 from typing_extensions import Protocol, Self
 from math import degrees, radians
 
+from OCP.TCollection import TCollection_HAsciiString
 from OCP.TDocStd import TDocStd_Document
 from OCP.TCollection import TCollection_ExtendedString
 from OCP.XCAFDoc import (
     XCAFDoc_DocumentTool,
     XCAFDoc_ColorType,
     XCAFDoc_ColorGen,
+    XCAFDoc_Material,
+    XCAFDoc_VisMaterial,
 )
 from OCP.XCAFApp import XCAFApp_Application
 from OCP.BinXCAFDrivers import BinXCAFDrivers
@@ -58,6 +61,129 @@
 AssemblyObjects = Union[Shape, Workplane, None]
 
 
+class Material(object):
+    """
+    Wrapper for the OCCT material classes XCAFDoc_Material and XCAFDoc_VisMaterial.
+    XCAFDoc_Material is focused on physical material properties and
+    XCAFDoc_VisMaterial is for visual properties to be used when rendering.
+    """
+
+    wrapped: XCAFDoc_Material
+    wrapped_vis: XCAFDoc_VisMaterial
+
+    def __init__(self, name: str | None = None, **kwargs):
+        """
+        Can be passed an arbitrary string name for the material along with keyword
+        arguments defining some other characteristics of the material. If nothing is
+        passed, arbitrary defaults are used.
+        """
+
+        # Create the default material object and prepare to set a few defaults
+        self.wrapped = XCAFDoc_Material()
+
+        # Default values in case the user did not set any others
+        aName = "Default"
+        aDescription = "Default material with properties similar to low carbon steel"
+        aDensity = 7.85
+        aDensityName = "Mass density"
+        aDensityTypeName = "g/cm^3"
+
+        # See if there are any non-defaults to be set
+        if name:
+            aName = name
+        if "description" in kwargs.keys():
+            aDescription = kwargs["description"]
+        if "density" in kwargs.keys():
+            aDensity = kwargs["density"]
+        if "densityUnit" in kwargs.keys():
+            aDensityTypeName = kwargs["densityUnit"]
+
+        # Set the properties on the material object
+        self.wrapped.Set(
+            TCollection_HAsciiString(aName),
+            TCollection_HAsciiString(aDescription),
+            aDensity,
+            TCollection_HAsciiString(aDensityName),
+            TCollection_HAsciiString(aDensityTypeName),
+        )
+
+        # Create the default visual material object and allow it to be used just with
+        # the OCC layer, for now. When this material class is expanded to include visual
+        # attributes, the OCC docs say that XCAFDoc_VisMaterialTool should be used to
+        # manage those attributes on the XCAFDoc_VisMaterial class.
+        self.wrapped_vis = XCAFDoc_VisMaterial()
+
+    @property
+    def name(self) -> str:
+        """
+        Get the string name of the material.
+        """
+        return self.wrapped.GetName().ToCString()
+
+    @property
+    def description(self) -> str:
+        """
+        Get the string description of the material.
+        """
+        return self.wrapped.GetDescription().ToCString()
+
+    @property
+    def density(self) -> float:
+        """
+        Get the density value of the material.
+        """
+        return self.wrapped.GetDensity()
+
+    @property
+    def densityUnit(self) -> str:
+        """
+        Get the units that the material density is defined in.
+        """
+        return self.wrapped.GetDensValType().ToCString()
+
+    def toTuple(self) -> Tuple[str, str, float, str]:
+        """
+        Convert Material to a tuple.
+        """
+        name = self.name
+        description = self.description
+        density = self.density
+        densityUnit = self.densityUnit
+
+        return (name, description, density, densityUnit)
+
+    def __hash__(self):
+        """
+        Create a unique hash for this material via its tuple.
+        """
+        return hash(self.toTuple())
+
+    def __eq__(self, other):
+        """
+        Check equality of this material against another via its tuple.
+        """
+        return self.toTuple() == other.toTuple()
+
+    def __getstate__(self) -> Tuple[str, str, float, str]:
+        """
+        Allows pickling.
+        """
+        return self.toTuple()
+
+    def __setstate__(self, data: Tuple[str, str, float, str]):
+        """
+        Allows pickling.
+        """
+        self.wrapped = XCAFDoc_Material()
+        self.wrapped.Set(
+            TCollection_HAsciiString(data[0]),
+            TCollection_HAsciiString(data[1]),
+            data[2],
+            TCollection_HAsciiString("Mass density"),
+            TCollection_HAsciiString(data[3]),
+        )
+
+
 class Color(object):
     """
     Wrapper for the OCCT color object Quantity_ColorRGBA.
@@ -239,6 +365,7 @@ def add(
         loc: Optional[Location] = None,
         name: Optional[str] = None,
         color: Optional[Color] = None,
+        material: Optional[Union[Material, str]] = None,
         metadata: Optional[Dict[str, Any]] = None,
     ) -> Self:
         ...
@@ -249,6 +376,7 @@ def add(
         loc: Optional[Location] = None,
         name: Optional[str] = None,
         color: Optional[Color] = None,
+        material: Optional[Union[Material, str]] = None,
         metadata: Optional[Dict[str, Any]] = None,
         **kwargs: Any,
     ) -> Self:
diff --git a/tests/test_assembly.py b/tests/test_assembly.py
@@ -36,6 +36,7 @@
 from OCP.TDF import TDF_ChildIterator
 from OCP.Quantity import Quantity_ColorRGBA, Quantity_TOC_sRGB, Quantity_NameOfColor
 from OCP.TopAbs import TopAbs_ShapeEnum
+from OCP.Graphic3d import Graphic3d_NameOfMaterial
 
 
 @pytest.fixture(scope="function")
@@ -1527,6 +1528,55 @@ def test_colors_assy1(assy_fixture, request, tmpdir, kind):
     check_assy(assy, assy_i)
 
 
+def test_materials():
+    # Test a default material not attached to an assembly
+    mat_0 = cq.Material()
+    assert mat_0.name == "Default"
+
+    # Simple objects to be added to the assembly with the material
+    wp_1 = cq.Workplane().box(10, 10, 10)
+    wp_2 = cq.Workplane().box(5, 5, 5)
+    wp_3 = cq.Workplane().box(2.5, 2.5, 2.5)
+
+    # Add the object to the assembly with the material
+    assy = cq.Assembly()
+
+    # Test with a default material
+    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.density == 7.85
+    assert assy.children[0].material.densityUnit == "g/cm^3"
+
+    # Test with a user-defined material when passing properties in constructor
+    mat_2 = cq.Material(
+        "test", description="Test material", density=1.0, densityUnit="lb/in^3"
+    )
+    assy.add(wp_2, material=mat_2)
+    assert assy.children[1].material.name == "test"
+    assert assy.children[1].material.description == "Test material"
+    assert assy.children[1].material.density == 1.0
+    assert assy.children[1].material.densityUnit == "lb/in^3"
+
+    # The visualization material is left for later expansion
+    assert assy.children[1].material.wrapped_vis.IsEmpty()
+
+    # 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
+    assert mat_2.__hash__() == hash(("test", "Test material", 1.0, "lb/in^3"))
+
+    # Test the equality operator with material
+    assert mat_2 == cq.Material(
+        "test", description="Test material", density=1.0, densityUnit="lb/in^3"
+    )
+
+
 @pytest.mark.parametrize(
     "assy_fixture, expected",
     [
diff --git a/tests/test_pickle.py b/tests/test_pickle.py
@@ -10,6 +10,7 @@
     Assembly,
     Color,
     Workplane,
+    Material,
 )
 from cadquery.func import box
 
@@ -42,6 +43,14 @@ def test_shape():
 
 def test_assy():
 
-    assy = Assembly().add(box(1, 1, 1), color=Color("blue")).add(box(2, 2, 2))
+    mat_1 = Material(
+        "test", description="Test material", density=1.0, densityUnit="lb/in^3"
+    )
+
+    assy = (
+        Assembly()
+        .add(box(1, 1, 1), color=Color("blue"), material=mat_1)
+        .add(box(2, 2, 2))
+    )
 
     assert isinstance(loads(dumps(assy)), Assembly)
