diff --git a/cadquery/assembly.py b/cadquery/assembly.py index db4a2e96c..647ab75bf 100644 --- a/cadquery/assembly.py +++ b/cadquery/assembly.py @@ -543,6 +543,28 @@ def _flatten(self, parents=[]): return rv + def __iter__( + self, + loc: Optional[Location] = None, + name: Optional[str] = None, + color: Optional[Color] = None, + ) -> Iterator[Tuple[Shape, str, Location, Optional[Color]]]: + """ + Assembly iterator yielding shapes, names, locations and colors. + """ + + name = f"{name}/{self.name}" if name else self.name + loc = loc * self.loc if loc else self.loc + color = self.color if self.color else color + + if self.obj: + yield self.obj if isinstance(self.obj, Shape) else Compound.makeCompound( + s for s in self.obj.vals() if isinstance(s, Shape) + ), name, loc, color + + for ch in self.children: + yield from ch.__iter__(loc, name, color) + def toCompound(self) -> Compound: """ Returns a Compound made from this Assembly (including all children) with the diff --git a/cadquery/occ_impl/assembly.py b/cadquery/occ_impl/assembly.py index 2143f6163..daefb22bc 100644 --- a/cadquery/occ_impl/assembly.py +++ b/cadquery/occ_impl/assembly.py @@ -1,4 +1,15 @@ -from typing import Union, Iterable, Tuple, Dict, overload, Optional, Any, List, cast +from typing import ( + Union, + Iterable, + Iterator, + Tuple, + Dict, + overload, + Optional, + Any, + List, + cast, +) from typing_extensions import Protocol from math import degrees @@ -12,7 +23,7 @@ from OCP.Quantity import Quantity_ColorRGBA from OCP.BRepAlgoAPI import BRepAlgoAPI_Fuse from OCP.TopTools import TopTools_ListOfShape -from OCP.BOPAlgo import BOPAlgo_GlueEnum +from OCP.BOPAlgo import BOPAlgo_GlueEnum, BOPAlgo_MakeConnected from OCP.TopoDS import TopoDS_Shape from vtkmodules.vtkRenderingCore import ( @@ -21,8 +32,11 @@ vtkRenderer, ) +from vtkmodules.vtkFiltersExtraction import vtkExtractCellsByType +from vtkmodules.vtkCommonDataModel import VTK_TRIANGLE, VTK_LINE, VTK_VERTEX + from .geom import Location -from .shapes import Shape, Compound +from .shapes import Shape, Solid, Compound from .exporters.vtk import toString from ..cq import Workplane @@ -131,6 +145,14 @@ def children(self) -> Iterable["AssemblyProtocol"]: def traverse(self) -> Iterable[Tuple[str, "AssemblyProtocol"]]: ... + def __iter__( + self, + loc: Optional[Location] = None, + name: Optional[str] = None, + color: Optional[Color] = None, + ) -> Iterator[Tuple[Shape, str, Location, Optional[Color]]]: + ... + def setName(l: TDF_Label, name: str, tool): @@ -227,45 +249,70 @@ def _toCAF(el, ancestor, color) -> TDF_Label: def toVTK( assy: AssemblyProtocol, - renderer: vtkRenderer = vtkRenderer(), - loc: Location = Location(), color: Tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0), tolerance: float = 1e-3, angularTolerance: float = 0.1, ) -> vtkRenderer: - loc = loc * assy.loc - trans, rot = loc.toTuple() + renderer = vtkRenderer() + + for shape, _, loc, col_ in assy: + + col = col_.toTuple() if col_ else color + trans, rot = loc.toTuple() + + data = shape.toVtkPolyData(tolerance, angularTolerance) + + # extract faces + extr = vtkExtractCellsByType() + extr.SetInputDataObject(data) + + extr.AddCellType(VTK_LINE) + extr.AddCellType(VTK_VERTEX) + extr.Update() + data_edges = extr.GetOutput() + + # extract edges + extr = vtkExtractCellsByType() + extr.SetInputDataObject(data) - if assy.color: - color = assy.color.toTuple() + extr.AddCellType(VTK_TRIANGLE) + extr.Update() + data_faces = extr.GetOutput() - if assy.shapes: - data = Compound.makeCompound(assy.shapes).toVtkPolyData( - tolerance, angularTolerance - ) + # remove normals from edges + data_edges.GetPointData().RemoveArray("Normals") + # add both to the renderer mapper = vtkMapper() - mapper.SetInputData(data) + mapper.AddInputDataObject(data_faces) actor = vtkActor() actor.SetMapper(mapper) actor.SetPosition(*trans) actor.SetOrientation(*map(degrees, rot)) - actor.GetProperty().SetColor(*color[:3]) - actor.GetProperty().SetOpacity(color[3]) + actor.GetProperty().SetColor(*col[:3]) + actor.GetProperty().SetOpacity(col[3]) renderer.AddActor(actor) - for child in assy.children: - renderer = toVTK(child, renderer, loc, color, tolerance, angularTolerance) + mapper = vtkMapper() + mapper.AddInputDataObject(data_edges) + + actor = vtkActor() + actor.SetMapper(mapper) + actor.SetPosition(*trans) + actor.SetOrientation(*map(degrees, rot)) + actor.GetProperty().SetColor(0, 0, 0) + actor.GetProperty().SetLineWidth(2) + + renderer.AddActor(actor) return renderer def toJSON( assy: AssemblyProtocol, - loc: Location = Location(), color: Tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0), tolerance: float = 1e-3, ) -> List[Dict[str, Any]]: @@ -273,29 +320,22 @@ def toJSON( Export an object to a structure suitable for converting to VTK.js JSON. """ - loc = loc * assy.loc - trans, rot = loc.toTuple() - - if assy.color: - color = assy.color.toTuple() - rv = [] - if assy.shapes: + for shape, _, loc, col_ in assy: + val: Any = {} data = toString(Compound.makeCompound(assy.shapes), tolerance) + trans, rot = loc.toTuple() val["shape"] = data - val["color"] = color + val["color"] = col_.toTuple() if col_ else color val["position"] = trans val["orientation"] = rot rv.append(val) - for child in assy.children: - rv.extend(toJSON(child, loc, color, tolerance)) - return rv @@ -331,19 +371,9 @@ def toFusedCAF( shapes: List[Shape] = [] colors = [] - def extract_shapes(assy, parent_loc=None, parent_color=None): - - loc = parent_loc * assy.loc if parent_loc else assy.loc - color = assy.color if assy.color else parent_color - - for shape in assy.shapes: - shapes.append(shape.moved(loc).copy()) - colors.append(color) - - for ch in assy.children: - extract_shapes(ch, loc, color) - - extract_shapes(assy) + for shape, _, loc, color in assy: + shapes.append(shape.moved(loc).copy()) + colors.append(color) # Initialize with a dummy value for mypy top_level_shape = cast(TopoDS_Shape, None) @@ -411,3 +441,37 @@ def extract_shapes(assy, parent_loc=None, parent_color=None): color_tool.SetColor(cur_lbl, color.wrapped, XCAFDoc_ColorGen) return top_level_lbl, doc + + +def imprint(assy: AssemblyProtocol) -> Tuple[Shape, Dict[Shape, Tuple[str, ...]]]: + """ + Imprint all the solids and construct a dictionary mapping imprinted solids to names from the input assy. + """ + + # make the id map + id_map = {} + + for obj, name, loc, _ in assy: + for s in obj.moved(loc).Solids(): + id_map[s] = name + + # connect topologically + bldr = BOPAlgo_MakeConnected() + bldr.SetRunParallel(True) + bldr.SetUseOBB(True) + + for obj in id_map: + bldr.AddArgument(obj.wrapped) + + bldr.Perform() + res = Shape(bldr.Shape()) + + # make the connected solid -> id map + origins: Dict[Shape, Tuple[str, ...]] = {} + + for s in res.Solids(): + ids = tuple(id_map[Solid(el)] for el in bldr.GetOrigins(s.wrapped)) + # if GetOrigins yields nothing, solid was not modified + origins[s] = ids if ids else (id_map[s],) + + return res, origins diff --git a/cadquery/occ_impl/exporters/assembly.py b/cadquery/occ_impl/exporters/assembly.py index 9a7ce9050..df28f00ad 100644 --- a/cadquery/occ_impl/exporters/assembly.py +++ b/cadquery/occ_impl/exporters/assembly.py @@ -140,10 +140,9 @@ def _vtkRenderWindow( Convert an assembly to a vtkRenderWindow. Used by vtk based exporters. """ - renderer = vtkRenderer() + renderer = toVTK(assy, tolerance=tolerance, angularTolerance=angularTolerance) renderWindow = vtkRenderWindow() renderWindow.AddRenderer(renderer) - toVTK(assy, renderer, tolerance=tolerance, angularTolerance=angularTolerance) renderer.ResetCamera() renderer.SetBackground(1, 1, 1) diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 21864be09..1c7199115 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -1475,3 +1475,42 @@ def test_point_constraint(simple_assy2): t2 = assy.children[1].loc.wrapped.Transformation().TranslationPart() assert t2.Modulus() == pytest.approx(1) + + +@pytest.fixture +def touching_assy(): + + b1 = cq.Workplane().box(1, 1, 1) + b2 = cq.Workplane(origin=(1, 0, 0)).box(1, 1, 1) + + return cq.Assembly().add(b1).add(b2) + + +@pytest.fixture +def disjoint_assy(): + + b1 = cq.Workplane().box(1, 1, 1) + b2 = cq.Workplane(origin=(2, 0, 0)).box(1, 1, 1) + + return cq.Assembly().add(b1).add(b2) + + +def test_imprinting(touching_assy, disjoint_assy): + + # normal usecase + r, o = cq.occ_impl.assembly.imprint(touching_assy) + + assert len(r.Solids()) == 2 + assert len(r.Faces()) == 11 + + for s in r.Solids(): + assert s in o + + # edge usecase + r, o = cq.occ_impl.assembly.imprint(disjoint_assy) + + assert len(r.Solids()) == 2 + assert len(r.Faces()) == 12 + + for s in r.Solids(): + assert s in o