Add STEP Import for Assemblies (#1779)

* Initial commit with colors and locations working for simple assemblies

* Trying to fix mypy error

* Jumping through hoops to make mypy happy

* Lint fix

* Refactored method interface based on suggestion

* Use Self and @classmethod

* Lint fix

* Added a test specifically for testing metadata

* Seeing if coverage and tests will pass

* Added name loading from STEP file

* Made name a public property of Assembly

* Trying to increase test coverage a bit

* Syncing up some experiments

* Got color and layer search working, still need to get name search working through indirect lookup

* Got tests working for layer and color info import

* Got shape name loading to work

* Trying to get approximate tuple comparison working

* Black fix

* Added a test for a bad filename, and added a custom color

* Increase test coverage a bit and improve color name check

* Removing code that should never be hit

* Still trying to increase test coverage

* Added a test for a plain assembly

* Refactored a bit to support nested assemblies better in the future

* Fixed location handling for components of the assembly

* Fix the default color to not be black

* Fixed bug with parent location not being applied when needed

* Fixed importer for Assembly.export method

* Fixed comment

* Removed a stray import that was probably added by AI somewhere along the line.

* Implement some of the suggestions

* Added a test for nested subassemblies on import

* Rework which covers everything except subshapes and layers

* Added layer name support back in

* Added a round-trip test and fixed issues that it revealed

* mypy fixes

* More mypy fixes

* Missed a cast

* More mypy fixes

* Tried to remove the attribute iterator and could not, but moved some code out of loop

* Fixes and simplifications based on codecov checks

* More cleanup to try to get code coverage high enough without creating a contrived STEP file

* Fix lack of application to the top-level assembly

* Trying to increase test coverage high enough

* Implemented changes based on PR from Adam

* Fix mypy errors

* Add isSubshape

* Add __getitem__

* Mypy fix and a usability tweak.

* Generic support for subshape metadata

* Rework assy import

* Start reworking tests

* Streamline a bit

* Add validation when adding subshapes

* Fix pickling

* Simplify assy  import

* Fix subshape name handling

* Simplify

* Additional test

* Better test

* Remove some changes

* Fix some warnings

* Fix top level key in .objects

* Initial commit

* mypy fix

* Check .objects keys in the test

* Black fix

* Rework

* Mypy fix

* Fix name

* Add xbf and start using xml xcaf

* Test with a simpler model

* Start on xml/xbf import

* Fix path handling

* Rework color handling

* Fix subshapes for xbf/xml

* Cleanup

* Remove additional levels and add/fix roundtrip tests

* Tweak tests

* Eliminate dead code

* Docs update

* Typo fixes + clarification

* typo fix

---------

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

Author: jmwright

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
 from .occ_impl.geom import Location
 from .occ_impl.assembly import Color
 from .occ_impl.solver import (
@@ -34,13 +35,15 @@
     exportGLTF,
     STEPExportModeLiterals,
 )
+from .occ_impl.importers.assembly import importStep as _importStep, importXbf, importXml
 
 from .selectors import _expression_grammar as _selector_grammar
 from .utils import deprecate
 
 # type definitions
 AssemblyObjects = Union[Shape, Workplane, None]
-ExportLiterals = Literal["STEP", "XML", "GLTF", "VTKJS", "VRML", "STL"]
+ImportLiterals = Literal["STEP", "XML", "XBF"]
+ExportLiterals = Literal["STEP", "XML", "XBF", "GLTF", "VTKJS", "VRML", "STL"]
 
 PATH_DELIM = "/"
 
@@ -155,6 +158,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 +179,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 +201,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.
 
@@ -344,11 +351,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
@@ -360,13 +367,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):
@@ -411,7 +418,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.
         """
@@ -506,7 +513,7 @@ def save(
         tolerance: float = 0.1,
         angularTolerance: float = 0.1,
         **kwargs,
-    ) -> "Assembly":
+    ) -> Self:
         """
         Save assembly to a file.
 
@@ -521,38 +528,9 @@ def save(
         :type ascii: bool
         """
 
-        # Make sure the export mode setting is correct
-        if mode not in get_args(STEPExportModeLiterals):
-            raise ValueError(f"Unknown assembly export mode {mode} for STEP")
-
-        if exportType is None:
-            t = path.split(".")[-1].upper()
-            if t in ("STEP", "XML", "VRML", "VTKJS", "GLTF", "GLB", "STL"):
-                exportType = cast(ExportLiterals, t)
-            else:
-                raise ValueError("Unknown extension, specify export type explicitly")
-
-        if exportType == "STEP":
-            exportAssembly(self, path, mode, **kwargs)
-        elif exportType == "XML":
-            exportCAF(self, path)
-        elif exportType == "VRML":
-            exportVRML(self, path, tolerance, angularTolerance)
-        elif exportType == "GLTF" or exportType == "GLB":
-            exportGLTF(self, path, None, tolerance, angularTolerance)
-        elif exportType == "VTKJS":
-            exportVTKJS(self, path)
-        elif exportType == "STL":
-            # Handle the ascii setting for STL export
-            export_ascii = False
-            if "ascii" in kwargs:
-                export_ascii = bool(kwargs.get("ascii"))
-
-            self.toCompound().exportStl(path, tolerance, angularTolerance, export_ascii)
-        else:
-            raise ValueError(f"Unknown format: {exportType}")
-
-        return self
+        return self.export(
+            path, exportType, mode, tolerance, angularTolerance, **kwargs
+        )
 
     def export(
         self,
@@ -562,7 +540,7 @@ def export(
         tolerance: float = 0.1,
         angularTolerance: float = 0.1,
         **kwargs,
-    ) -> "Assembly":
+    ) -> Self:
         """
         Save assembly to a file.
 
@@ -583,7 +561,7 @@ def export(
 
         if exportType is None:
             t = path.split(".")[-1].upper()
-            if t in ("STEP", "XML", "VRML", "VTKJS", "GLTF", "GLB", "STL"):
+            if t in ("STEP", "XML", "XBF", "VRML", "VTKJS", "GLTF", "GLB", "STL"):
                 exportType = cast(ExportLiterals, t)
             else:
                 raise ValueError("Unknown extension, specify export type explicitly")
@@ -592,6 +570,8 @@ def export(
             exportAssembly(self, path, mode, **kwargs)
         elif exportType == "XML":
             exportCAF(self, path)
+        elif exportType == "XBF":
+            exportCAF(self, path, binary=True)
         elif exportType == "VRML":
             exportVRML(self, path, tolerance, angularTolerance)
         elif exportType == "GLTF" or exportType == "GLB":
@@ -611,9 +591,39 @@ def export(
         return self
 
     @classmethod
-    def load(cls, path: str) -> "Assembly":
+    def importStep(cls, path: str) -> Self:
+        """
+        Reads an assembly from a STEP file.
 
-        raise NotImplementedError
+        :param path: Path and filename for reading.
+        :return: An Assembly object.
+        """
+
+        return cls.load(path, importType="STEP")
+
+    @classmethod
+    def load(cls, path: str, importType: Optional[ImportLiterals] = None,) -> Self:
+        """
+        Load step, xbf or xml.
+        """
+
+        if importType is None:
+            t = path.split(".")[-1].upper()
+            if t in ("STEP", "XML", "XBF"):
+                importType = cast(ImportLiterals, t)
+            else:
+                raise ValueError("Unknown extension, specify export type explicitly")
+
+        assy = cls()
+
+        if importType == "STEP":
+            _importStep(assy, path)
+        elif importType == "XML":
+            importXml(assy, path)
+        elif importType == "XBF":
+            importXbf(assy, path)
+
+        return assy
 
     @property
     def shapes(self) -> List[Shape]:
@@ -714,12 +724,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,
 )