From b51715abd3ede06f08451e013ccab321f50400be Mon Sep 17 00:00:00 2001
From: Romain FERRU <romain.ferru@gmail.com>
Date: Wed, 23 Feb 2022 22:54:43 +0100
Subject: [PATCH] Topology traversal selection

minimal working example
---
 cadquery/cq.py                     |  26 ++++++-
 cadquery/occ_impl/topo_explorer.py | 105 +++++++++++++++++++++++++++++
 2 files changed, 130 insertions(+), 1 deletion(-)
 create mode 100644 cadquery/occ_impl/topo_explorer.py

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,
+    )