Assembly hierarchy - 2nd take (#545)

* Reworked nested assy querying * Support subclassing * Added docstring
CadQuery · Dec 18, 2020 · 82a043a · 82a043a
1 parent a576f1d
commit 82a043a
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 10 deletions.
diff --git a/cadquery/assembly.py b/cadquery/assembly.py
@@ -21,15 +21,26 @@
 ConstraintKinds = Literal["Plane", "Point", "Axis"]
 ExportLiterals = Literal["STEP", "XML"]
 
+PATH_DELIM = "/"
+
 # enitity selector grammar definiiton
 def _define_grammar():
 
-    from pyparsing import Literal as Literal, Word, Optional, alphas, alphanums
+    from pyparsing import (
+        Literal as Literal,
+        Word,
+        Optional,
+        alphas,
+        alphanums,
+        delimitedList,
+    )
 
     Separator = Literal("@").suppress()
     TagSeparator = Literal("?").suppress()
 
-    Name = Word(alphas, alphanums + "_").setResultsName("name")
+    Name = delimitedList(
+        Word(alphas, alphanums + "_"), PATH_DELIM, combine=True
+    ).setResultsName("name")
     Tag = Word(alphas, alphanums + "_").setResultsName("tag")
     Selector = _selector_grammar.setResultsName("selector")
 
@@ -234,6 +245,11 @@ def add(self, arg, **kwargs):
 
         if isinstance(arg, Assembly):
 
+            # enforce unique names
+            name = kwargs["name"] if kwargs.get("name") else arg.name
+            if name in self.objects:
+                raise ValueError("Unique name is required")
+
             subassy = arg._copy()
 
             subassy.loc = kwargs["loc"] if kwargs.get("loc") else arg.loc
@@ -242,11 +258,10 @@ def add(self, arg, **kwargs):
             subassy.parent = self
 
             self.children.append(subassy)
-            self.objects[subassy.name] = subassy
-            self.objects.update(subassy.objects)
+            self.objects.update(subassy._flatten())
 
         else:
-            assy = Assembly(arg, **kwargs)
+            assy = self.__class__(arg, **kwargs)
             assy.parent = self
 
             self.add(assy)
@@ -454,3 +469,17 @@ def traverse(self) -> Iterator[Tuple[str, "Assembly"]]:
                 yield el
 
         yield (self.name, self)
+
+    def _flatten(self, parents=[]):
+        """
+        Generate a dict with all ancestors with keys indicating parent-child relations.
+        """
+
+        rv = {}
+
+        for ch in self.children:
+            rv.update(ch._flatten(parents=parents + [self.name]))
+
+        rv[PATH_DELIM.join(parents + [self.name])] = self
+
+        return rv
diff --git a/tests/test_assembly.py b/tests/test_assembly.py
@@ -153,8 +153,8 @@ def test_constrain(simple_assy, nested_assy):
 
     assert len(simple_assy.constraints) == 3
 
-    nested_assy.constrain("TOP@faces@>Z", "BOTTOM@faces@<Z", "Plane")
-    nested_assy.constrain("TOP@faces@>X", "BOTTOM@faces@<X", "Axis")
+    nested_assy.constrain("TOP@faces@>Z", "SECOND/BOTTOM@faces@<Z", "Plane")
+    nested_assy.constrain("TOP@faces@>X", "SECOND/BOTTOM@faces@<X", "Axis")
 
     assert len(nested_assy.constraints) == 2
 
@@ -168,7 +168,7 @@ def test_constrain(simple_assy, nested_assy):
         .IsEqual(gp_XYZ(), 1e-9)
     )
     assert constraint.sublocs[1].wrapped.IsEqual(
-        nested_assy.objects["BOTTOM"].loc.wrapped
+        nested_assy.objects["SECOND/BOTTOM"].loc.wrapped
     )
 
     simple_assy.solve()
@@ -199,10 +199,16 @@ def test_constrain(simple_assy, nested_assy):
 def test_constrain_with_tags(nested_assy):
 
     nested_assy.add(None, name="dummy")
-    nested_assy.constrain("TOP?top_face", "BOTTOM", "Plane")
+    nested_assy.constrain("TOP?top_face", "SECOND/BOTTOM", "Plane")
 
     assert len(nested_assy.constraints) == 1
 
     # test selection of a non-shape object
     with pytest.raises(ValueError):
-        nested_assy.constrain("BOTTOM ? pts", "dummy", "Plane")
+        nested_assy.constrain("SECOND/BOTTOM ? pts", "dummy", "Plane")
+
+
+def test_duplicate_name(nested_assy):
+
+    with pytest.raises(ValueError):
+        nested_assy.add(None, name="SECOND")