Flatter assy export

cadquery/assembly.py

@@ -11,12 +11,13 @@
     cast,
     get_args,
 )
-from typing_extensions import Literal
+from typing_extensions import Literal, Self
 from typish import instance_of
 from uuid import uuid1 as uuid
+from warnings import warn
 
 from .cq import Workplane
-from .occ_impl.shapes import Shape, Compound
+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 (
@@ -34,6 +35,7 @@
     exportGLTF,
     STEPExportModeLiterals,
 )
+from .occ_impl.importers.assembly import importStep as _importStep
 
 from .selectors import _expression_grammar as _selector_grammar
 from .utils import deprecate
@@ -155,6 +157,10 @@ 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)
+
         for ch in self.children:
             ch_copy = ch._copy()
             ch_copy.parent = rv
@@ -172,7 +178,7 @@ def add(
         loc: Optional[Location] = None,
         name: Optional[str] = None,
         color: Optional[Color] = None,
-    ) -> "Assembly":
+    ) -> Self:
         """
         Add a subassembly to the current assembly.
 
@@ -194,7 +200,7 @@ def add(
         name: Optional[str] = None,
         color: Optional[Color] = None,
         metadata: Optional[Dict[str, Any]] = None,
-    ) -> "Assembly":
+    ) -> Self:
         """
         Add a subassembly to the current assembly with explicit location and name.
 
@@ -342,11 +348,11 @@ def _subloc(self, name: str) -> Tuple[Location, str]:
     @overload
     def constrain(
         self, q1: str, q2: str, kind: ConstraintKind, param: Any = None
-    ) -> "Assembly":
+    ) -> Self:
         ...
 
     @overload
-    def constrain(self, q1: str, kind: ConstraintKind, param: Any = None) -> "Assembly":
+    def constrain(self, q1: str, kind: ConstraintKind, param: Any = None) -> Self:
         ...
 
     @overload
@@ -358,13 +364,13 @@ def constrain(
         s2: Shape,
         kind: ConstraintKind,
         param: Any = None,
-    ) -> "Assembly":
+    ) -> Self:
         ...
 
     @overload
     def constrain(
         self, id1: str, s1: Shape, kind: ConstraintKind, param: Any = None,
-    ) -> "Assembly":
+    ) -> Self:
         ...
 
     def constrain(self, *args, param=None):
@@ -409,7 +415,7 @@ def constrain(self, *args, param=None):
 
         return self
 
-    def solve(self, verbosity: int = 0) -> "Assembly":
+    def solve(self, verbosity: int = 0) -> Self:
         """
         Solve the constraints.
         """
@@ -504,7 +510,7 @@ def save(
         tolerance: float = 0.1,
         angularTolerance: float = 0.1,
         **kwargs,
-    ) -> "Assembly":
+    ) -> Self:
         """
         Save assembly to a file.
 
@@ -560,7 +566,7 @@ def export(
         tolerance: float = 0.1,
         angularTolerance: float = 0.1,
         **kwargs,
-    ) -> "Assembly":
+    ) -> Self:
         """
         Save assembly to a file.
 
@@ -609,9 +615,26 @@ def export(
         return self
 
     @classmethod
-    def load(cls, path: str) -> "Assembly":
+    def importStep(cls, path: str) -> Self:
+        """
+        Reads an assembly from a STEP file.
+
+        :param path: Path and filename for writing.
+        :return: An Assembly object.
+        """
+
+        assy = cls()
+        _importStep(assy, path)
+
+        return assy
+
+    @classmethod
+    def load(cls, path: str) -> Self:
+        """
+        Alias of importStep for now.
+        """
 
-        raise NotImplementedError
+        return cls.importStep(path)
 
     @property
     def shapes(self) -> List[Shape]:
@@ -712,12 +735,81 @@ def addSubshape(
         :return: The modified assembly.
         """
 
+        # check if the subshape belongs to the stored object
+        if any(isSubshape(s, obj) for obj in self.shapes):
+            assy = self
+        else:
+            warn(
+                "Current node does not contain any Shapes, searching in subnodes. In the future this will result in an error."
+            )
+
+            found = False
+            for ch in self.children:
+                if any(isSubshape(s, obj) for obj in ch.shapes):
+                    assy = ch
+                    found = True
+                    break
+
+            if not found:
+                raise ValueError(
+                    f"{s} is not a subshape of the current node or its children"
+                )
+
         # Handle any metadata we were passed
         if name:
-            self._subshape_names[s] = name
+            assy._subshape_names[s] = name
         if color:
-            self._subshape_colors[s] = color
+            assy._subshape_colors[s] = color
         if layer:
-            self._subshape_layers[s] = layer
+            assy._subshape_layers[s] = layer
 
         return self
+
+    def __getitem__(self, name: str) -> "Assembly":
+        """
+        [] based access to children.
+        """
+
+        return self.objects[name]
+
+    def _ipython_key_completions_(self) -> List[str]:
+        """
+        IPython autocompletion helper.
+        """
+
+        return list(self.objects.keys())
+
+    def __contains__(self, name: str) -> bool:
+
+        return name in self.objects
+
+    def __getattr__(self, name: str) -> "Assembly":
+        """
+        . based access to children.
+        """
+
+        if name in self.objects:
+            return self.objects[name]
+
+        raise AttributeError
+
+    def __dir__(self):
+        """
+        Modified __dir__ for autocompletion.
+        """
+
+        return list(self.__dict__) + list(ch.name for ch in self.children)
+
+    def __getstate__(self):
+        """
+        Explicit getstate needed due to getattr.
+        """
+
+        return self.__dict__
+
+    def __setstate__(self, d):
+        """
+        Explicit setstate needed due to getattr.
+        """
+
+        self.__dict__ = d

cadquery/func.py

@@ -51,4 +51,5 @@
     setThreads,
     project,
     faceOn,
+    isSubshape,
 )