Skip to content

Commit 045b2fb

Browse files
adam-urbanczykadam-urbanczyk
committed
Initial commit
1 parent 6e1651c commit 045b2fb

File tree

4 files changed

+89
-63
lines changed

4 files changed

+89
-63
lines changed

‎cadquery/occ_impl/assembly.py‎

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -299,9 +299,10 @@ def toCAF(
299299

300300
def _toCAF(el, ancestor, color) -> TDF_Label:
301301

302-
# create a subassy
303-
subassy = tool.NewShape()
304-
setName(subassy, el.name, tool)
302+
# create a subassy if needed
303+
if el.children:
304+
subassy = tool.NewShape()
305+
setName(subassy, el.name, tool)
305306

306307
# define the current color
307308
current_color = el.color if el.color else color
@@ -326,7 +327,7 @@ def _toCAF(el, ancestor, color) -> TDF_Label:
326327
compounds[key1] = compound
327328

328329
tool.SetShape(lab, compound.wrapped)
329-
setName(lab, f"{el.name}_part", tool)
330+
setName(lab, f"{el.name}_part" if el.children else el.name, tool)
330331
unique_objs[key0] = lab
331332

332333
# handle colors when exporting to STEP
@@ -366,19 +367,30 @@ def _toCAF(el, ancestor, color) -> TDF_Label:
366367
)
367368
ltool.SetLayer(subshape_label, layer_label)
368369

369-
tool.AddComponent(subassy, lab, TopLoc_Location())
370+
if el.children:
371+
lab = tool.AddComponent(subassy, lab, TopLoc_Location())
372+
setName(lab, f"{el.name}_part", tool)
373+
else:
374+
lab = tool.AddComponent(ancestor, lab, el.loc.wrapped)
375+
setName(lab, f"{el.name}", tool)
370376

371377
# handle colors when *not* exporting to STEP
372378
if not coloredSTEP and current_color:
373-
setColor(subassy, current_color, ctool)
379+
if el.children:
380+
setColor(subassy, current_color, ctool)
381+
382+
if el.obj:
383+
setColor(lab, current_color, ctool)
374384

375385
# add children recursively
376386
for child in el.children:
377387
_toCAF(child, subassy, current_color)
378388

379-
if ancestor:
389+
if ancestor and el.children:
380390
tool.AddComponent(ancestor, subassy, el.loc.wrapped)
381391
rv = subassy
392+
elif ancestor:
393+
rv = ancestor
382394
else:
383395
# update the top level location
384396
rv = TDF_Label() # NB: additional label is needed to apply the location

‎cadquery/occ_impl/exporters/assembly.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -176,9 +176,12 @@ def _process_child(child: AssemblyProtocol, assy_label: TDF_Label):
176176
# Handle shape name, color and location
177177
part_label = shape_tool.AddShape(shape.wrapped, False)
178178
TDataStd_Name.Set_s(part_label, TCollection_ExtendedString(name))
179+
179180
if color:
180181
color_tool.SetColor(part_label, color.wrapped, XCAFDoc_ColorGen)
181-
shape_tool.AddComponent(assy_label, part_label, loc.wrapped)
182+
183+
ref_label = shape_tool.AddComponent(assy_label, part_label, loc.wrapped)
184+
TDataStd_Name.Set_s(ref_label, TCollection_ExtendedString(name))
182185

183186
# If this assembly has shape metadata, add it to the shape
184187
if (

‎cadquery/occ_impl/importers/assembly.py‎

Lines changed: 27 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,20 @@
1515
from ..shapes import Shape
1616

1717

18+
def _name(label: TDF_Label) -> str | None:
19+
"""
20+
Internal helper for getting names
21+
"""
22+
23+
rv = None
24+
25+
name_attr = TDataStd_Name()
26+
if label.FindAttribute(TDataStd_Name.GetID_s(), name_attr):
27+
rv = str(name_attr.Get().ToExtString())
28+
29+
return rv
30+
31+
1832
def importStep(assy: AssemblyProtocol, path: str):
1933
"""
2034
Import a step file into an assembly.
@@ -45,13 +59,13 @@ def _process_label(lbl: TDF_Label, parent: AssemblyProtocol):
4559
ref_label = TDF_Label()
4660
shape_tool.GetReferredShape_s(comp_label, ref_label)
4761

48-
# Find the name of this referenced part
49-
ref_name_attr = TDataStd_Name()
50-
if ref_label.FindAttribute(TDataStd_Name.GetID_s(), ref_name_attr):
51-
ref_name = str(ref_name_attr.Get().ToExtString())
52-
5362
if shape_tool.IsAssembly_s(ref_label):
5463

64+
# Find the name of this referenced part
65+
ref_name_attr = TDataStd_Name()
66+
if ref_label.FindAttribute(TDataStd_Name.GetID_s(), ref_name_attr):
67+
ref_name = str(ref_name_attr.Get().ToExtString())
68+
5569
sub_assy = assy.__class__(name=ref_name)
5670

5771
# Recursively process subassemblies
@@ -61,6 +75,10 @@ def _process_label(lbl: TDF_Label, parent: AssemblyProtocol):
6175
parent.add(sub_assy, name=ref_name, loc=cq_loc)
6276

6377
elif shape_tool.IsSimpleShape_s(ref_label):
78+
# Find the name of this referenced part
79+
ref_name_attr = TDataStd_Name()
80+
if comp_label.FindAttribute(TDataStd_Name.GetID_s(), ref_name_attr):
81+
ref_name = str(ref_name_attr.Get().ToExtString())
6482

6583
# A single shape needs to be added to the assembly
6684
final_shape = shape_tool.GetShape_s(ref_label)
@@ -184,7 +202,11 @@ def _process_label(lbl: TDF_Label, parent: AssemblyProtocol):
184202
# Set the name of the top-level assembly to match the top-level label
185203
name_attr = TDataStd_Name()
186204
top_level_label.FindAttribute(TDataStd_Name.GetID_s(), name_attr)
205+
206+
# Manipulation of .objects is needed to maintain consistency
207+
assy.objects.pop(assy.name)
187208
assy.name = str(name_attr.Get().ToExtString())
209+
assy.objects[assy.name] = assy
188210

189211
# Get the location of the top-level component
190212
comp_labels = TDF_LabelSequence()

‎tests/test_assembly.py‎

Lines changed: 39 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,8 @@ def empty_top_assy():
100100

101101
b1 = cq.Workplane().box(1, 1, 1)
102102

103-
assy = cq.Assembly()
104-
assy.add(b1, color=cq.Color("green"))
103+
assy = cq.Assembly(name="top")
104+
assy.add(b1, color=cq.Color("green"), name="b")
105105

106106
return assy
107107

@@ -441,10 +441,15 @@ def get_doc_nodes(doc, leaf=False):
441441
ctool = expl.ColorTool()
442442
style = node.Style
443443
label = node.RefLabel
444+
label2 = node.Label
444445

445446
name_att = TDataStd_Name()
446447
label.FindAttribute(TDataStd_Name.GetID_s(), name_att)
447448

449+
if label2.IsAttribute(TDataStd_Name.GetID_s()):
450+
name_att = TDataStd_Name()
451+
label2.FindAttribute(TDataStd_Name.GetID_s(), name_att)
452+
448453
color = style.GetColorSurfRGBA()
449454
shape = expl.FindShapeFromPathId_s(doc, node.Id)
450455
color_shape = Quantity_ColorRGBA()
@@ -1082,9 +1087,12 @@ def make_model(name: str, COPY: bool):
10821087
assy = Assembly(name="test_assy")
10831088
assy.add(box(1, 2, 5), color=Color("green"))
10841089

1085-
for v in rect(10, 10).vertices():
1090+
for i, v in enumerate(rect(10, 10).vertices()):
10861091
assy.add(
1087-
b.copy() if COPY else b, loc=Location(v.Center()), color=Color("red")
1092+
b.copy() if COPY else b,
1093+
name=f"element_{i}",
1094+
loc=Location(v.Center()),
1095+
color=Color("red"),
10881096
)
10891097

10901098
assy.export(name)
@@ -1417,27 +1425,11 @@ def test_leaf_node_count(assy_fixture, count, request):
14171425
[
14181426
(
14191427
["chassis", "wheel-axle.*", "wheel:.*"],
1420-
{
1421-
"color": (1.0, 0.0, 0.0, 1.0),
1422-
"color_shape": (1.0, 0.0, 0.0, 1.0),
1423-
"num_nodes": 4,
1424-
},
1425-
),
1426-
(
1427-
["chassis", "wheel-axle.*", "wheel:.*", "wheel.*_part"],
1428-
{"color": (1.0, 0.0, 0.0, 1.0), "num_nodes": 4},
1428+
{"color": (1.0, 0.0, 0.0, 1.0), "num_nodes": 4,},
14291429
),
14301430
(
14311431
["chassis", "wheel-axle.*", "axle"],
1432-
{
1433-
"color": (0.0, 1.0, 0.0, 1.0),
1434-
"color_shape": (0.0, 1.0, 0.0, 1.0),
1435-
"num_nodes": 2,
1436-
},
1437-
),
1438-
(
1439-
["chassis", "wheel-axle.*", "axle", "axle_part"],
1440-
{"color": (0.0, 1.0, 0.0, 1.0), "num_nodes": 2},
1432+
{"color": (0.0, 1.0, 0.0, 1.0), "num_nodes": 2,},
14411433
),
14421434
],
14431435
),
@@ -1478,103 +1470,100 @@ def check_nodes(doc, expected):
14781470
{"color_shape": (0.0, 1.0, 0.0, 1.0)},
14791471
),
14801472
(
1481-
["TOP", "SECOND", "BOTTOM", "BOTTOM_part"],
1473+
["TOP", "SECOND", "BOTTOM"],
14821474
{
14831475
"color_shape": (0.0, 1.0, 0.0, 1.0),
14841476
"color_subshapes": (0.0, 1.0, 0.0, 1.0),
14851477
},
14861478
),
14871479
],
14881480
),
1489-
("empty_top_assy", [([".*_part"], {"color_shape": (0.0, 1.0, 0.0, 1.0)}),]),
1481+
("empty_top_assy", [(["top", "b"], {"color_shape": (0.0, 1.0, 0.0, 1.0)}),]),
14901482
(
14911483
"boxes0_assy",
14921484
[
1493-
(["box0", "box0_part"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1494-
(["box1", "box0_part"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1485+
(["box0"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1486+
(["box1"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
14951487
],
14961488
),
14971489
(
14981490
"boxes1_assy",
14991491
[
1500-
(["box0", "box0_part"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1501-
(["box1", "box0_part"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1492+
(["box0"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1493+
(["box1"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
15021494
],
15031495
),
15041496
(
15051497
"boxes2_assy",
15061498
[
1507-
(["box0", "box0_part"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1508-
(["box1", "box1_part"], {"color_shape": (0.0, 1.0, 0.0, 1.0)}),
1499+
(["box0"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1500+
(["box1"], {"color_shape": (0.0, 1.0, 0.0, 1.0)}),
15091501
],
15101502
),
15111503
(
15121504
"boxes3_assy",
15131505
[
1514-
(["box0", "box0_part"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1506+
(["box0"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
15151507
(
1516-
["box1", "box1_part"],
1508+
["box1"],
15171509
{"color_shape": cq.Color().toTuple()},
15181510
), # default color when unspecified
15191511
],
15201512
),
15211513
(
15221514
"boxes4_assy",
15231515
[
1524-
(["box_0", "box_0_part"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1525-
(["box_1", "box_1_part"], {"color_shape": (0.0, 1.0, 0.0, 1.0)}),
1516+
(["box_0"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1517+
(["box_1"], {"color_shape": (0.0, 1.0, 0.0, 1.0)}),
15261518
],
15271519
),
15281520
(
15291521
"boxes5_assy",
15301522
[
1531-
(["box:a", "box:a_part"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1532-
(["box:b", "box:b_part"], {"color_shape": (0.0, 1.0, 0.0, 1.0)}),
1523+
(["box:a"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1524+
(["box:b"], {"color_shape": (0.0, 1.0, 0.0, 1.0)}),
15331525
],
15341526
),
15351527
(
15361528
"boxes6_assy",
15371529
[
1538-
(["box__0", "box__0_part"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1539-
(["box__1", "box__1_part"], {"color_shape": (0.0, 1.0, 0.0, 1.0)}),
1530+
(["box__0"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1531+
(["box__1"], {"color_shape": (0.0, 1.0, 0.0, 1.0)}),
15401532
],
15411533
),
15421534
(
15431535
"boxes7_assy",
15441536
[
1545-
(["box_0", "box_0_part"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1546-
(["box", "box_part"], {"color_shape": (0.0, 1.0, 0.0, 1.0)}),
1547-
(
1548-
["another box", "another box_part"],
1549-
{"color_shape": (0.23, 0.26, 0.26, 0.6)},
1550-
),
1537+
(["box_0"], {"color_shape": (1.0, 0.0, 0.0, 1.0)}),
1538+
(["box"], {"color_shape": (0.0, 1.0, 0.0, 1.0)}),
1539+
(["another box"], {"color_shape": (0.23, 0.26, 0.26, 0.6)},),
15511540
],
15521541
),
15531542
(
15541543
"chassis0_assy",
15551544
[
15561545
(
1557-
["chassis", "wheel-axle-front", "wheel:left", "wheel:left_part"],
1546+
["chassis", "wheel-axle-front", "wheel:left"],
15581547
{"color_shape": (1.0, 0.0, 0.0, 1.0)},
15591548
),
15601549
(
1561-
["chassis", "wheel-axle-front", "wheel:right", "wheel:left_part"],
1550+
["chassis", "wheel-axle-front", "wheel:right"],
15621551
{"color_shape": (1.0, 0.0, 0.0, 1.0)},
15631552
),
15641553
(
1565-
["chassis", "wheel-axle-rear", "wheel:left", "wheel:left_part"],
1554+
["chassis", "wheel-axle-rear", "wheel:left"],
15661555
{"color_shape": (1.0, 0.0, 0.0, 1.0)},
15671556
),
15681557
(
1569-
["chassis", "wheel-axle-rear", "wheel:right", "wheel:left_part"],
1558+
["chassis", "wheel-axle-rear", "wheel:right"],
15701559
{"color_shape": (1.0, 0.0, 0.0, 1.0)},
15711560
),
15721561
(
1573-
["chassis", "wheel-axle-front", "axle", "axle_part"],
1562+
["chassis", "wheel-axle-front", "axle"],
15741563
{"color_shape": (0.0, 1.0, 0.0, 1.0)},
15751564
),
15761565
(
1577-
["chassis", "wheel-axle-rear", "axle", "axle_part"],
1566+
["chassis", "wheel-axle-rear", "axle"],
15781567
{"color_shape": (0.0, 1.0, 0.0, 1.0)},
15791568
),
15801569
],

0 commit comments

Comments
 (0)