Fix new assembly subshape referencing functionality

Author: adam-urbanczyk

cadquery/assembly.py

@@ -47,6 +47,8 @@
 PATH_DELIM = "/"
 
 # entity selector grammar definition
+
+
 def _define_grammar():
 
     from pyparsing import (
@@ -765,10 +767,10 @@ def addSubshape(
 
         return self
 
-    def __getitem__(self, name: str) -> "Assembly":
+    def __getitem__(self, name: str) -> Union["Assembly", Shape]:
         """
         [] based access to children.
-        
+
         """
 
         if name in self.objects:
@@ -777,7 +779,7 @@ def __getitem__(self, name: str) -> "Assembly":
             rv = self.objects[name[0]]
 
             if name[1] in rv._subshape_names.inv:
-                rv = self._subshape_names.inv[name[1]]
+                rv = compound(self._subshape_names.inv[name[1]])
                 return rv[0] if len(rv) == 1 else compound(rv)
 
         raise KeyError
@@ -793,7 +795,7 @@ def __contains__(self, name: str) -> bool:
 
         return name in self.objects
 
-    def __getattr__(self, name: str) -> "Assembly":
+    def __getattr__(self, name: str) -> Union["Assembly", Shape]:
         """
         . based access to children.
         """

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

@@ -10,6 +10,8 @@
 from OCP.XCAFDoc import XCAFDoc_ColorSurf, XCAFDoc_DocumentTool, XCAFDoc_GraphNode
 from OCP.Interface import Interface_Static
 
+from typing import cast
+
 from ..assembly import AssemblyProtocol, Color
 from ..geom import Location
 from ..shapes import Shape
@@ -94,7 +96,7 @@ def _process_label(lbl: TDF_Label, parent: AssemblyProtocol):
                         parent.add(tmp)
 
                         # change the current assy to handle subshape data
-                        current = parent[ref_name]
+                        current = cast(AssemblyProtocol, parent[ref_name])
 
                     # iterate over subshape and handle names, layers and colors
                     subshape_labels = TDF_LabelSequence()
@@ -206,7 +208,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])
 
         # Copy all of the children over to the main assembly object
         for child in imported_assy.children:

cadquery/utils.py

@@ -95,7 +95,7 @@ class BiDict(UserDict[K, V]):
     Bi-directional dictionary.
     """
 
-    _inv = dict[V, list[K]]
+    _inv: dict[V, list[K]]
 
     def __init__(self, *args, **kwargs):