Enable bidirectional mapping of subshapes (#1949)

* Add assy changes

* Extend special methods to handle subshape names

---------

Co-authored-by: adam-urbanczyk <13981538+adam-urbanczyk@users.noreply.github.com>

Author: adam-urbanczyk

cadquery/assembly.py

@@ -15,9 +15,10 @@
 from typish import instance_of
 from uuid import uuid1 as uuid
 from warnings import warn
+from itertools import chain
 
 from .cq import Workplane
-from .occ_impl.shapes import Shape, Compound, isSubshape
+from .occ_impl.shapes import Shape, Compound, isSubshape, compound
 from .occ_impl.geom import Location
 from .occ_impl.assembly import Color
 from .occ_impl.solver import (
@@ -38,7 +39,7 @@
 from .occ_impl.importers.assembly import importStep as _importStep, importXbf, importXml
 
 from .selectors import _expression_grammar as _selector_grammar
-from .utils import deprecate
+from .utils import deprecate, BiDict
 
 # type definitions
 AssemblyObjects = Union[Shape, Workplane, None]
@@ -48,6 +49,8 @@
 PATH_DELIM = "/"
 
 # entity selector grammar definition
+
+
 def _define_grammar():
 
     from pyparsing import (
@@ -98,9 +101,9 @@ class Assembly(object):
     constraints: List[Constraint]
 
     # Allows metadata to be stored for exports
-    _subshape_names: dict[Shape, str]
-    _subshape_colors: dict[Shape, Color]
-    _subshape_layers: dict[Shape, str]
+    _subshape_names: BiDict[Shape, str]
+    _subshape_colors: BiDict[Shape, Color]
+    _subshape_layers: BiDict[Shape, str]
 
     _solve_result: Optional[Dict[str, Any]]
 
@@ -147,9 +150,9 @@ def __init__(
 
         self._solve_result = None
 
-        self._subshape_names = {}
-        self._subshape_colors = {}
-        self._subshape_layers = {}
+        self._subshape_names = BiDict()
+        self._subshape_colors = BiDict()
+        self._subshape_layers = BiDict()
 
     def _copy(self) -> "Assembly":
         """
@@ -158,9 +161,9 @@ def _copy(self) -> "Assembly":
 
         rv = self.__class__(self.obj, self.loc, self.name, self.color, self.metadata)
 
-        rv._subshape_colors = dict(self._subshape_colors)
-        rv._subshape_names = dict(self._subshape_names)
-        rv._subshape_layers = dict(self._subshape_layers)
+        rv._subshape_colors = BiDict(self._subshape_colors)
+        rv._subshape_names = BiDict(self._subshape_names)
+        rv._subshape_layers = BiDict(self._subshape_layers)
 
         for ch in self.children:
             ch_copy = ch._copy()
@@ -754,40 +757,54 @@ def addSubshape(
 
         return self
 
-    def __getitem__(self, name: str) -> "Assembly":
+    def __getitem__(self, name: str) -> Union["Assembly", Shape]:
         """
         [] based access to children.
+
         """
 
-        return self.objects[name]
+        if name in self.objects:
+            return self.objects[name]
+        elif name in self._subshape_names.inv:
+            rv = self._subshape_names.inv[name]
+            return rv[0] if len(rv) == 1 else compound(rv)
+
+        raise KeyError
 
     def _ipython_key_completions_(self) -> List[str]:
         """
         IPython autocompletion helper.
         """
 
-        return list(self.objects.keys())
+        return list(chain(self.objects.keys(), self._subshape_names.inv.keys()))
 
     def __contains__(self, name: str) -> bool:
 
-        return name in self.objects
+        return name in self.objects or name in self._subshape_names.inv
 
-    def __getattr__(self, name: str) -> "Assembly":
+    def __getattr__(self, name: str) -> Union["Assembly", Shape]:
         """
         . based access to children.
         """
 
         if name in self.objects:
             return self.objects[name]
+        elif name in self._subshape_names.inv:
+            rv = self._subshape_names.inv[name]
+            return rv[0] if len(rv) == 1 else compound(rv)
 
-        raise AttributeError
+        raise AttributeError(f"{name} is not an attribute of {self}")
 
     def __dir__(self):
         """
         Modified __dir__ for autocompletion.
         """
 
-        return list(self.__dict__) + list(ch.name for ch in self.children)
+        return (
+            list(self.__dict__)
+            + list(ch.name for ch in self.children)
+            + list(self._subshape_names.inv.keys())
+        )
 
     def __getstate__(self):
         """

cadquery/occ_impl/assembly.py

@@ -52,6 +52,7 @@
 from .shapes import Shape, Solid, Compound
 from .exporters.vtk import toString
 from ..cq import Workplane
+from ..utils import BiDict
 
 # type definitions
 AssemblyObjects = Union[Shape, Workplane, None]
@@ -210,15 +211,15 @@ def children(self) -> Iterable["AssemblyProtocol"]:
         ...
 
     @property
-    def _subshape_names(self) -> Dict[Shape, str]:
+    def _subshape_names(self) -> BiDict[Shape, str]:
         ...
 
     @property
-    def _subshape_colors(self) -> Dict[Shape, Color]:
+    def _subshape_colors(self) -> BiDict[Shape, Color]:
         ...
 
     @property
-    def _subshape_layers(self) -> Dict[Shape, str]:
+    def _subshape_layers(self) -> BiDict[Shape, str]:
         ...
 
     @overload
@@ -276,7 +277,7 @@ def __iter__(
     ) -> Iterator[Tuple[Shape, str, Location, Optional[Color]]]:
         ...
 
-    def __getitem__(self, name: str) -> Self:
+    def __getitem__(self, name: str) -> Self | Shape:
         ...
 
     def __contains__(self, name: str) -> bool:

cadquery/occ_impl/importers/assembly.py

@@ -253,7 +253,7 @@ def _process_label(lbl: TDF_Label, parent: AssemblyProtocol):
                         parent.add(tmp)
 
                         # change the current assy to handle subshape data
-                        current = parent[comp_name]
+                        current = cast(AssemblyProtocol, parent[comp_name])
 
                     # iterate over subshape and handle names, layers and colors
                     subshape_labels = TDF_LabelSequence()
@@ -326,6 +326,7 @@ def _process_label(lbl: TDF_Label, parent: AssemblyProtocol):
         assy.objects.pop(assy.name)
         assy.name = str(name_attr.Get().ToExtString())
         assy.objects[assy.name] = assy
+
         if cq_color:
             assy.color = cq_color
 
@@ -350,7 +351,7 @@ def _process_label(lbl: TDF_Label, parent: AssemblyProtocol):
         # extras on successive round-trips. exportStepMeta does not add the extra top-level
         # node and so does not exhibit this behavior.
         if assy.name in imported_assy:
-            imported_assy = imported_assy[assy.name]
+            imported_assy = cast(AssemblyProtocol, imported_assy[assy.name])
             # comp_labels = TDF_LabelSequence()
             # shape_tool.GetComponents_s(top_level_label, comp_labels)
             # comp_label = comp_labels.Value(1)

cadquery/utils.py

@@ -2,6 +2,7 @@
 from inspect import signature, isbuiltin
 from typing import TypeVar, Callable, cast
 from warnings import warn
+from collections import UserDict
 
 from multimethod import multimethod, DispatchError
 
@@ -83,3 +84,35 @@ def get_arity(f: TCallable) -> int:
         rv = f.__code__.co_argcount - n_defaults
 
     return rv
+
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class BiDict(UserDict[K, V]):
+    """
+    Bi-directional dictionary.
+    """
+
+    _inv: dict[V, list[K]]
+
+    def __init__(self, *args, **kwargs):
+
+        self._inv = {}
+
+        super().__init__(*args, **kwargs)
+
+    def __setitem__(self, k: K, v: V):
+
+        super().__setitem__(k, v)
+
+        if v in self._inv:
+            self._inv[v].append(k)
+        else:
+            self._inv[v] = [k]
+
+    @property
+    def inv(self) -> dict[V, list[K]]:
+
+        return self._inv

tests/test_assembly.py

@@ -17,7 +17,7 @@
     exportVRML,
 )
 from cadquery.occ_impl.assembly import toJSON, toCAF, toFusedCAF
-from cadquery.occ_impl.shapes import Face, box, cone, plane
+from cadquery.occ_impl.shapes import Face, box, cone, plane, Compound
 
 from OCP.gp import gp_XYZ
 from OCP.TDocStd import TDocStd_Document
@@ -407,6 +407,10 @@ def subshape_assy():
         layer="cylinder_bottom_wire_layer",
     )
 
+    # Add two subshapes with the same name
+    assy["cyl_1"].addSubshape(cyl_1.faces(">Z").val(), name="2_faces")
+    assy["cyl_1"].addSubshape(cyl_1.faces("<Z").val(), name="2_faces")
+
     return assy
 
 
@@ -2330,6 +2334,20 @@ def test_special_methods(subshape_assy):
         subshape_assy.cube_123456
 
 
+def test_subshape_access(subshape_assy):
+    """
+    Smoke-test subshape access.
+    """
+
+    assert "cube_1_top_face" in subshape_assy.cube_1.__dir__()
+    assert "cube_1_top_face" in subshape_assy.cube_1._ipython_key_completions_()
+    assert "cube_1_top_face" in subshape_assy.cube_1
+
+    assert isinstance(subshape_assy.cube_1.cube_1_top_face, Face)
+    assert isinstance(subshape_assy.cube_1["cube_1_top_face"], Face)
+    assert isinstance(subshape_assy.cyl_1["2_faces"], Compound)
+
+
 def test_shallow_assy():
     """
     toCAF edge case.