diff --git a/cadquery/cq.py b/cadquery/cq.py index cbbe65bde..e10a6bd97 100644 --- a/cadquery/cq.py +++ b/cadquery/cq.py @@ -37,11 +37,14 @@ from typing_extensions import Literal from inspect import Parameter, Signature +from cadquery.occ_impl.topo_explorer import ConnectedShapesExplorer + from .occ_impl.geom import Vector, Plane, Location -from .occ_impl.shapes import ( +from cadquery.occ_impl.shapes import ( Shape, Edge, + Shapes, Wire, Face, Solid, @@ -1055,6 +1058,27 @@ def compounds( """ return self._selectObjects("Compounds", selector, tag) + def connected(self, shape_type: Shapes, include_child_shape=False): + """ + Returns a new Workplane with shapes of type `shape_type` that are connected to the previous Workplane objects + Example: + code:: + box = Workplane().box(10,10,10).faces(">Z").connected("Edge") + print(len(box.val())) + >>> 4 # It selects the 4 edges orthogonal to the face + + Note that by default it doesn't include searched shapes from the selected child shape, if you wanted + the 4 edges orthogonal to the face as well as the 4 edges from the face set `include_child_shape` to True + """ + + if len(self.objects) > 1: + raise NotImplementedError("Only one object workplane supported yet") + explorer = ConnectedShapesExplorer(self.findSolid(), self.val()) + + objs = explorer.search(shape_type, include_child_shape) + + return self.newObject(objs) + def toSvg(self, opts: Any = None) -> str: """ Returns svg text that represents the first item on the stack. diff --git a/cadquery/occ_impl/topo_explorer.py b/cadquery/occ_impl/topo_explorer.py new file mode 100644 index 000000000..5cce2aa09 --- /dev/null +++ b/cadquery/occ_impl/topo_explorer.py @@ -0,0 +1,105 @@ +from typing import Optional +from OCP.TopExp import TopExp_Explorer + + +from OCP.TopAbs import ( + TopAbs_COMPOUND, + TopAbs_COMPSOLID, + TopAbs_SOLID, + TopAbs_SHELL, + TopAbs_FACE, + TopAbs_WIRE, + TopAbs_EDGE, + TopAbs_VERTEX, + TopAbs_SHAPE, +) +from cadquery.occ_impl.shapes import ( + Shape, + Shapes, +) + + +ENUM_MAPPING = { + "Compound": TopAbs_COMPOUND, + "CompSolid": TopAbs_COMPSOLID, + "Solid": TopAbs_SOLID, + "Shell": TopAbs_SHELL, + "Face": TopAbs_FACE, + "Wire": TopAbs_WIRE, + "Edge": TopAbs_EDGE, + "Vertex": TopAbs_VERTEX, + "Shape": TopAbs_SHAPE, +} + + +class ShapeExplorer: + def __init__(self, shape: Shape) -> None: + self.shape = shape.wrapped + self.explorer = TopExp_Explorer() + + def search( + self, shape_type: Shapes, not_from: Optional[Shapes] = None, + ): + """ + Searchs all the shapes of type `shape_type` in the shape. If `not_from` is specified, will avoid all the shapes + that are attached to the type `not_from` + """ + to_avoid = TopAbs_SHAPE if not_from is None else ENUM_MAPPING[not_from] + self.explorer.Init(self.shape, ENUM_MAPPING[shape_type], to_avoid) + + collection = [] + while self.explorer.More(): + shape = Shape.cast(self.explorer.Current()) + collection.append(shape) + self.explorer.Next() + + return list(set(collection)) # the 'set' is used to remove duplicates + + +class ConnectedShapesExplorer: + def __init__(self, base_shape, child_shape) -> None: + self.base_shape = base_shape + self.child_shape = child_shape + self.explorer = ShapeExplorer(base_shape) + + def _connected_by_vertices(self, shape, by_all=False): + child_vertices = self.child_shape.Vertices() + shape_vertices = shape.Vertices() + + if by_all: + return all(v in child_vertices for v in shape_vertices) + else: + return any(v in child_vertices for v in shape_vertices) + + def search(self, shape_type: Shapes, include_child_shape=False): + candidate_shapes = self.explorer.search(shape_type) + if not include_child_shape: + child_shapes = ShapeExplorer(self.child_shape).search(shape_type) + candidate_shapes = [ + shape for shape in candidate_shapes if shape not in child_shapes + ] + + connected_shapes = [] + for shape in candidate_shapes: + if self._connected_by_vertices(shape): + connected_shapes.append(shape) + return connected_shapes + + +if __name__ == "__main__": + import cadquery as cq + from jupyter_cadquery.viewer.client import show + + box = cq.Workplane().box(10, 10, 10).faces(">Z").connected("Edge", True) + + show( + box, + height=800, + cad_width=1500, + reset_camera=False, + default_edgecolor=(255, 255, 255), + zoom=1, + axes=True, + axes0=True, + render_edges=True, + )