Assembly support

cadquery/__init__.py

@@ -28,13 +28,15 @@
     Selector,
 )
 from .cq import CQ, Workplane
+from .assembly import Assembly
 from . import selectors
 from . import plugins
 
 
 __all__ = [
     "CQ",
     "Workplane",
+    "Assembly",
     "plugins",
     "selectors",
     "Plane",
@@ -64,4 +66,4 @@
     "plugins",
 ]
 
-__version__ = "2.0"
+__version__ = "2.1dev"

cadquery/assembly.py

@@ -0,0 +1,202 @@
+from typing import Union, Optional, List, Mapping, Any, overload, Tuple, Iterator, cast
+from typing_extensions import Literal
+from uuid import uuid1 as uuid
+
+from .cq import Workplane
+from .occ_impl.shapes import Shape
+from .occ_impl.geom import Location
+from .occ_impl.exporters.assembly import exportAssembly, exportCAF
+
+
+AssemblyObjects = Union[Shape, Workplane, None]
+ConstraintKinds = Literal["Plane", "Point", "Axis"]
+ExportLiterals = Literal["STEP", "XML"]
+
+
+class Constraint(object):
+
+    objects: Tuple[Shape, ...]
+    args: Tuple[Shape, ...]
+    kind: ConstraintKinds
+
+
+class Assembly(object):
+    """Nested assembly of Workplane and Shape objects defining their relative positions.
+    """
+
+    loc: Location
+    name: str
+    metadata: Mapping[str, Any]
+
+    obj: AssemblyObjects
+    parent: Optional["Assembly"]
+    children: List["Assembly"]
+
+    objects: Mapping[str, AssemblyObjects]
+    constraints: List[Constraint]
+
+    def __init__(
+        self,
+        obj: AssemblyObjects = None,
+        loc: Optional[Location] = None,
+        name: Optional[str] = None,
+    ):
+        """
+        construct an assembly
+
+        :param obj: root object of the assembly (deafault: None)
+        :param loc: location of the root object (deafault: None, interpreted as identity transformation)
+        :param name: unique name of the root object (default: None, reasulting in an UUID being generated)
+        :return: An Assembly object.
+
+        To create an empt assembly use::
+
+            assy = Assembly(None)
+
+        To create one containt a root object::
+
+            b = Workplane().box(1,1,1)
+            assy = Assembly(b, Location(Vector(0,0,1)), name="root")
+
+        """
+
+        self.obj = obj
+        self.loc = loc if loc else Location()
+        self.name = name if name else str(uuid())
+        self.parent = None
+
+        self.children = []
+        self.objects = {self.name: self.obj}
+
+    @overload
+    def add(
+        self,
+        obj: "Assembly",
+        loc: Optional[Location] = None,
+        name: Optional[str] = None,
+    ) -> "Assembly":
+        """
+        add a subassembly to the current assembly.
+
+        :param obj: subassembly to be added
+        :param loc: location of the root object (deafault: None, resulting in the location stored in the subassembly being used)
+        :param name: unique name of the root object (default: None, resulting in the name stored in the subassembly being used)
+        """
+        ...
+
+    @overload
+    def add(
+        self,
+        obj: AssemblyObjects,
+        loc: Optional[Location] = None,
+        name: Optional[str] = None,
+    ) -> "Assembly":
+        """
+        add a subassembly to the current assembly with explicit location and name
+
+        :param obj: object to be added as a subassembly
+        :param loc: location of the root object (deafault: None, interpreted as identity transformation)
+        :param name: unique name of the root object (default: None, resulting in an UUID being generated)
+        """
+        ...
+
+    def add(self, arg, **kwargs):
+
+        if isinstance(arg, Assembly):
+
+            subassy = Assembly(
+                arg.obj,
+                kwargs["loc"] if kwargs.get("loc") else arg.loc,
+                kwargs["name"] if kwargs.get("name") else arg.name,
+            )
+
+            subassy.children.extend(arg.children)
+            subassy.objects[subassy.name] = subassy.obj
+            subassy.objects.update(arg.objects)
+
+            self.children.append(subassy)
+            self.objects[subassy.name] = subassy.obj
+            self.objects.update(subassy.objects)
+
+            arg.parent = self
+
+        else:
+            assy = Assembly(arg, **kwargs)
+            assy.parent = self
+
+            self.add(assy)
+
+        return self
+
+    @overload
+    def constrain(self, query: str) -> "Assembly":
+        ...
+
+    @overload
+    def constrain(self, q1: str, q2: str, kind: ConstraintKinds) -> "Assembly":
+        ...
+
+    @overload
+    def constrain(
+        self, id1: str, s1: Shape, id2: str, s2: Shape, kind: ConstraintKinds
+    ) -> "Assembly":
+        ...
+
+    def constrain(self, *args):
+
+        raise NotImplementedError
+
+    def solve(self) -> "Assembly":
+
+        raise NotImplementedError
+
+    def save(
+        self, path: str, exportType: Optional[ExportLiterals] = None
+    ) -> "Assembly":
+        """
+        save as STEP or OCCT native XML file
+
+        :param path: filepath
+        :param exportType: export format (deafault: None, results in format being inferred form the path)
+        """
+
+        if exportType is None:
+            t = path.split(".")[-1].upper()
+            if t in ("STEP", "XML"):
+                exportType = cast(ExportLiterals, t)
+            else:
+                raise ValueError("Unknown extension, specify export type explicitly")
+
+        if exportType == "STEP":
+            exportAssembly(self, path)
+        elif exportType == "XML":
+            exportCAF(self, path)
+        else:
+            raise ValueError(f"Unknown format: {exportType}")
+
+        return self
+
+    @classmethod
+    def load(cls, path: str) -> "Assembly":
+
+        raise NotImplementedError
+
+    @property
+    def shapes(self) -> List[Shape]:
+
+        rv: List[Shape] = []
+
+        if isinstance(self.obj, Shape):
+            rv = [self.obj]
+        elif isinstance(self.obj, Workplane):
+            rv = [el for el in self.obj.vals() if isinstance(el, Shape)]
+
+        return rv
+
+    def traverse(self) -> Iterator[Tuple[str, "Assembly"]]:
+
+        for ch in self.children:
+            for el in ch.traverse():
+                yield el
+
+        yield (self.name, self)

cadquery/occ_impl/assembly.py

@@ -0,0 +1,75 @@
+from typing import Iterable, Tuple, Dict
+from typing_extensions import Protocol
+
+from OCP.TDocStd import TDocStd_Document
+from OCP.TCollection import TCollection_ExtendedString
+from OCP.XCAFDoc import XCAFDoc_DocumentTool
+from OCP.TDataStd import TDataStd_Name
+from OCP.TDF import TDF_Label
+from OCP.TopLoc import TopLoc_Location
+
+from .geom import Location
+from .shapes import Shape, Compound
+
+
+class AssemblyProtocol(Protocol):
+    @property
+    def loc(self) -> Location:
+        ...
+
+    @property
+    def name(self) -> str:
+        ...
+
+    @property
+    def shapes(self) -> Iterable[Shape]:
+        ...
+
+    @property
+    def children(self) -> Iterable["AssemblyProtocol"]:
+        ...
+
+    def traverse(self) -> Iterable[Tuple[str, "AssemblyProtocol"]]:
+        ...
+
+
+def setName(l: TDF_Label, name, tool):
+
+    TDataStd_Name.Set_s(l, TCollection_ExtendedString(name))
+
+
+def toCAF(assy: AssemblyProtocol) -> Tuple[TDF_Label, TDocStd_Document]:
+
+    # prepare a doc
+    doc = TDocStd_Document(TCollection_ExtendedString("XmlOcaf"))
+    tool = XCAFDoc_DocumentTool.ShapeTool_s(doc.Main())
+    tool.SetAutoNaming_s(False)
+
+    # add root
+    top = tool.NewShape()
+    TDataStd_Name.Set_s(top, TCollection_ExtendedString("CQ assembly"))
+
+    # add leafs and subassemblies
+    subassys: Dict[str, Tuple[TDF_Label, Location]] = {}
+    for k, v in assy.traverse():
+        # leaf part
+        lab = tool.NewShape()
+        tool.SetShape(lab, Compound.makeCompound(v.shapes).wrapped)
+        setName(lab, f"{k}_part", tool)
+
+        # assy part
+        subassy = tool.NewShape()
+        tool.AddComponent(subassy, lab, TopLoc_Location())
+        setName(subassy, k, tool)
+
+        subassys[k] = (subassy, v.loc)
+
+        for ch in v.children:
+            tool.AddComponent(
+                subassy, subassys[ch.name][0], subassys[ch.name][1].wrapped
+            )
+
+    tool.AddComponent(top, subassys[assy.name][0], assy.loc.wrapped)
+    tool.UpdateAssemblies()
+
+    return top, doc

cadquery/occ_impl/exporters/assembly.py

@@ -0,0 +1,52 @@
+from OCP.XSControl import XSControl_WorkSession
+from OCP.STEPCAFControl import STEPCAFControl_Writer
+from OCP.STEPControl import STEPControl_StepModelType
+from OCP.IFSelect import IFSelect_ReturnStatus
+from OCP.XCAFApp import XCAFApp_Application
+from OCP.XmlDrivers import (
+    XmlDrivers_DocumentStorageDriver,
+    XmlDrivers_DocumentRetrievalDriver,
+)
+from OCP.TCollection import TCollection_ExtendedString, TCollection_AsciiString
+from OCP.PCDM import PCDM_StoreStatus
+
+from ..assembly import AssemblyProtocol, toCAF
+
+
+def exportAssembly(assy: AssemblyProtocol, path: str) -> bool:
+
+    _, doc = toCAF(assy)
+
+    session = XSControl_WorkSession()
+    writer = STEPCAFControl_Writer(session, False)
+    writer.SetNameMode(True)
+    writer.Transfer(doc, STEPControl_StepModelType.STEPControl_AsIs)
+
+    status = writer.Write(path)
+
+    return status == IFSelect_ReturnStatus.IFSelect_RetDone
+
+
+def exportCAF(assy: AssemblyProtocol, path: str) -> bool:
+
+    _, doc = toCAF(assy)
+    app = XCAFApp_Application.GetApplication_s()
+
+    store = XmlDrivers_DocumentStorageDriver(
+        TCollection_ExtendedString("Copyright: Open Cascade, 2001-2002")
+    )
+    ret = XmlDrivers_DocumentRetrievalDriver()
+
+    app.DefineFormat(
+        TCollection_AsciiString("XmlOcaf"),
+        TCollection_AsciiString("Xml XCAF Document"),
+        TCollection_AsciiString(path.split(".")[-1]),
+        ret,
+        store,
+    )
+
+    status = app.SaveAs(doc, TCollection_ExtendedString(path))
+
+    app.Close(doc)
+
+    return status == PCDM_StoreStatus.PCDM_SS_OK

cadquery/occ_impl/geom.py

@@ -862,7 +862,7 @@ def isInside(self, b2: "BoundBox") -> bool:
 
 class Location(object):
     """Location in 3D space. Depending on usage can be absolute or relative.
-    
+
     This class wraps the TopLoc_Location class from OCCT. It can be used to move Shape
     objects in both relative and absolute manner. It is the preferred type to locate objects
     in CQ.
@@ -895,6 +895,11 @@ def __init__(self, t: TopLoc_Location) -> None:
         """Location wrapping the low-level TopLoc_Location object t"""
         ...
 
+    @overload
+    def __init__(self, t: gp_Trsf) -> None:
+        """Location wrapping the low-level gp_Trsf object t"""
+        ...
+
     @overload
     def __init__(self, t: Vector, ax: Vector, angle: float) -> None:
         """Location with translation t and rotation around ax by angle
@@ -919,6 +924,8 @@ def __init__(self, *args):
             elif isinstance(t, TopLoc_Location):
                 self.wrapped = t
                 return
+            elif isinstance(t, gp_Trsf):
+                T = t
         elif len(args) == 2:
             t, v = args
             cs = gp_Ax3(v.toPnt(), t.zDir.toDir(), t.xDir.toDir())