-
Notifications
You must be signed in to change notification settings - Fork 276
/
assembly.py
594 lines (471 loc) · 17.5 KB
/
assembly.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
from functools import reduce
from typing import (
Union,
Optional,
List,
Dict,
Any,
overload,
Tuple,
Iterator,
cast,
get_args,
)
from typing_extensions import Literal
from typish import instance_of
from uuid import uuid1 as uuid
from .cq import Workplane
from .occ_impl.shapes import Shape, Compound
from .occ_impl.geom import Location
from .occ_impl.assembly import Color
from .occ_impl.solver import (
ConstraintKind,
ConstraintSolver,
ConstraintSpec as Constraint,
UnaryConstraintKind,
BinaryConstraintKind,
)
from .occ_impl.exporters.assembly import (
exportAssembly,
exportCAF,
exportVTKJS,
exportVRML,
exportGLTF,
STEPExportModeLiterals,
ExportModes,
)
from .selectors import _expression_grammar as _selector_grammar
# type definitions
AssemblyObjects = Union[Shape, Workplane, None]
ExportLiterals = Literal["STEP", "XML", "GLTF", "VTKJS", "VRML", "STL"]
PATH_DELIM = "/"
# entity selector grammar definition
def _define_grammar():
from pyparsing import (
Literal as Literal,
Word,
Optional,
alphas,
alphanums,
delimitedList,
)
Separator = Literal("@").suppress()
TagSeparator = Literal("?").suppress()
Name = delimitedList(
Word(alphas, alphanums + "_"), PATH_DELIM, combine=True
).setResultsName("name")
Tag = Word(alphas, alphanums + "_").setResultsName("tag")
Selector = _selector_grammar.setResultsName("selector")
SelectorType = (
Literal("solids") | Literal("faces") | Literal("edges") | Literal("vertices")
).setResultsName("selector_kind")
return (
Name
+ Optional(TagSeparator + Tag)
+ Optional(Separator + SelectorType + Separator + Selector)
)
_grammar = _define_grammar()
class Assembly(object):
"""Nested assembly of Workplane and Shape objects defining their relative positions."""
loc: Location
name: str
color: Optional[Color]
metadata: Dict[str, Any]
obj: AssemblyObjects
parent: Optional["Assembly"]
children: List["Assembly"]
objects: Dict[str, "Assembly"]
constraints: List[Constraint]
_solve_result: Optional[Dict[str, Any]]
def __init__(
self,
obj: AssemblyObjects = None,
loc: Optional[Location] = None,
name: Optional[str] = None,
color: Optional[Color] = None,
metadata: Optional[Dict[str, Any]] = None,
):
"""
construct an assembly
:param obj: root object of the assembly (default: None)
:param loc: location of the root object (default: None, interpreted as identity transformation)
:param name: unique name of the root object (default: None, resulting in an UUID being generated)
:param color: color of the added object (default: None)
:param metadata: a store for user-defined metadata (default: None)
:return: An Assembly object.
To create an empty assembly use::
assy = Assembly(None)
To create one constraint a root object::
b = Workplane().box(1, 1, 1)
assy = Assembly(b, Location(Vector(0, 0, 1)), name="root")
"""
self.obj = obj
self.loc = loc if loc else Location()
self.name = name if name else str(uuid())
self.color = color if color else None
self.metadata = metadata if metadata else {}
self.parent = None
self.children = []
self.constraints = []
self.objects = {self.name: self}
self._solve_result = None
def _copy(self) -> "Assembly":
"""
Make a deep copy of an assembly
"""
rv = self.__class__(self.obj, self.loc, self.name, self.color, self.metadata)
for ch in self.children:
ch_copy = ch._copy()
ch_copy.parent = rv
rv.children.append(ch_copy)
rv.objects[ch_copy.name] = ch_copy
rv.objects.update(ch_copy.objects)
return rv
@overload
def add(
self,
obj: "Assembly",
loc: Optional[Location] = None,
name: Optional[str] = None,
color: Optional[Color] = None,
) -> "Assembly":
"""
Add a subassembly to the current assembly.
:param obj: subassembly to be added
:param loc: location of the root object (default: None, resulting in the location stored in
the subassembly being used)
:param name: unique name of the root object (default: None, resulting in the name stored in
the subassembly being used)
:param color: color of the added object (default: None, resulting in the color stored in the
subassembly being used)
"""
...
@overload
def add(
self,
obj: AssemblyObjects,
loc: Optional[Location] = None,
name: Optional[str] = None,
color: Optional[Color] = None,
metadata: Optional[Dict[str, Any]] = None,
) -> "Assembly":
"""
Add a subassembly to the current assembly with explicit location and name.
:param obj: object to be added as a subassembly
:param loc: location of the root object (default: None, interpreted as identity
transformation)
:param name: unique name of the root object (default: None, resulting in an UUID being
generated)
:param color: color of the added object (default: None)
:param metadata: a store for user-defined metadata (default: None)
"""
...
def add(self, arg, **kwargs):
"""
Add a subassembly to the current assembly.
"""
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
subassy.name = kwargs["name"] if kwargs.get("name") else arg.name
subassy.color = kwargs["color"] if kwargs.get("color") else arg.color
subassy.metadata = (
kwargs["metadata"] if kwargs.get("metadata") else arg.metadata
)
subassy.parent = self
self.children.append(subassy)
self.objects.update(subassy._flatten())
else:
assy = self.__class__(arg, **kwargs)
assy.parent = self
self.add(assy)
return self
def _query(self, q: str) -> Tuple[str, Optional[Shape]]:
"""
Execute a selector query on the assembly.
The query is expected to be in the following format:
name[?tag][@kind@args]
valid example include:
obj_name @ faces @ >Z
obj_name?tag1@faces@>Z
obj_name ? tag
obj_name
"""
tmp: Workplane
res: Workplane
query = _grammar.parseString(q, True)
name: str = query.name
obj = self.objects[name].obj
if isinstance(obj, Workplane) and query.tag:
tmp = obj._getTagged(query.tag)
elif isinstance(obj, (Workplane, Shape)):
tmp = Workplane().add(obj)
else:
raise ValueError("Workplane or Shape required to define a constraint")
if query.selector:
res = getattr(tmp, query.selector_kind)(query.selector)
else:
res = tmp
val = res.val()
return name, val if isinstance(val, Shape) else None
def _subloc(self, name: str) -> Tuple[Location, str]:
"""
Calculate relative location of an object in a subassembly.
Returns the relative positions as well as the name of the top assembly.
"""
rv = Location()
obj = self.objects[name]
name_out = name
if obj not in self.children and obj is not self:
locs = []
while not obj.parent is self:
locs.append(obj.loc)
obj = cast(Assembly, obj.parent)
name_out = obj.name
rv = reduce(lambda l1, l2: l1 * l2, locs)
return (rv, name_out)
@overload
def constrain(
self, q1: str, q2: str, kind: ConstraintKind, param: Any = None
) -> "Assembly":
...
@overload
def constrain(self, q1: str, kind: ConstraintKind, param: Any = None) -> "Assembly":
...
@overload
def constrain(
self,
id1: str,
s1: Shape,
id2: str,
s2: Shape,
kind: ConstraintKind,
param: Any = None,
) -> "Assembly":
...
@overload
def constrain(
self, id1: str, s1: Shape, kind: ConstraintKind, param: Any = None,
) -> "Assembly":
...
def constrain(self, *args, param=None):
"""
Define a new constraint.
"""
# dispatch on arguments
if len(args) == 2:
q1, kind = args
id1, s1 = self._query(q1)
elif len(args) == 3 and instance_of(args[1], UnaryConstraintKind):
q1, kind, param = args
id1, s1 = self._query(q1)
elif len(args) == 3:
q1, q2, kind = args
id1, s1 = self._query(q1)
id2, s2 = self._query(q2)
elif len(args) == 4:
q1, q2, kind, param = args
id1, s1 = self._query(q1)
id2, s2 = self._query(q2)
elif len(args) == 5:
id1, s1, id2, s2, kind = args
elif len(args) == 6:
id1, s1, id2, s2, kind, param = args
else:
raise ValueError(f"Incompatible arguments: {args}")
# handle unary and binary constraints
if instance_of(kind, UnaryConstraintKind):
loc1, id1_top = self._subloc(id1)
c = Constraint((id1_top,), (s1,), (loc1,), kind, param)
elif instance_of(kind, BinaryConstraintKind):
loc1, id1_top = self._subloc(id1)
loc2, id2_top = self._subloc(id2)
c = Constraint((id1_top, id2_top), (s1, s2), (loc1, loc2), kind, param)
else:
raise ValueError(f"Unknown constraint: {kind}")
self.constraints.append(c)
return self
def solve(self, verbosity: int = 0) -> "Assembly":
"""
Solve the constraints.
"""
# Get all entities and number them. First entity is marked as locked
ents = {}
i = 0
locked: List[int] = []
for c in self.constraints:
for name in c.objects:
if name not in ents:
ents[name] = i
i += 1
if (c.kind == "Fixed" or name == self.name) and ents[
name
] not in locked:
locked.append(ents[name])
# Lock the first occurring entity if needed.
if not locked:
unary_objects = [
c.objects[0]
for c in self.constraints
if instance_of(c.kind, UnaryConstraintKind)
]
binary_objects = [
c.objects[0]
for c in self.constraints
if instance_of(c.kind, BinaryConstraintKind)
]
for b in binary_objects:
if b not in unary_objects:
locked.append(ents[b])
break
# Lock the first occurring entity if needed.
if not locked:
locked.append(0)
locs = [self.objects[n].loc for n in ents]
# construct the constraint mapping
constraints = []
for c in self.constraints:
ixs = tuple(ents[obj] for obj in c.objects)
pods = c.toPODs()
for pod in pods:
constraints.append((ixs, pod))
# check if any constraints were specified
if not constraints:
raise ValueError("At least one constraint required")
# check if at least two entities are present
if len(ents) < 2:
raise ValueError("At least two entities need to be constrained")
# instantiate the solver
scale = self.toCompound().BoundingBox().DiagonalLength
solver = ConstraintSolver(locs, constraints, locked=locked, scale=scale)
# solve
locs_new, self._solve_result = solver.solve(verbosity)
# update positions
# find the inverse root loc
loc_root_inv = Location()
if self.obj:
for loc_new, n in zip(locs_new, ents):
if n == self.name:
loc_root_inv = loc_new.inverse
break
# update the positions
for loc_new, n in zip(locs_new, ents):
if n != self.name:
self.objects[n].loc = loc_root_inv * loc_new
return self
def save(
self,
path: str,
exportType: Optional[ExportLiterals] = None,
mode: STEPExportModeLiterals = "default",
tolerance: float = 0.1,
angularTolerance: float = 0.1,
**kwargs,
) -> "Assembly":
"""
Save assembly to a file.
:param path: Path and filename for writing.
:param exportType: export format (default: None, results in format being inferred form the path)
:param mode: STEP only - See :meth:`~cadquery.occ_impl.exporters.assembly.exportAssembly`.
:param tolerance: the deflection tolerance, in model units. Only used for glTF, VRML. Default 0.1.
:param angularTolerance: the angular tolerance, in radians. Only used for glTF, VRML. Default 0.1.
:param \\**kwargs: Additional keyword arguments. Only used for STEP, glTF and STL.
See :meth:`~cadquery.occ_impl.exporters.assembly.exportAssembly`.
:param ascii: STL only - Sets whether or not STL export should be text or binary
: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
@classmethod
def load(cls, path: str) -> "Assembly":
raise NotImplementedError
@property
def shapes(self) -> List[Shape]:
"""
List of Shape objects in the .obj field
"""
rv: List[Shape] = []
if isinstance(self.obj, Shape):
rv = [self.obj]
elif isinstance(self.obj, Workplane):
rv = [el for el in self.obj.vals() if isinstance(el, Shape)]
return rv
def traverse(self) -> Iterator[Tuple[str, "Assembly"]]:
"""
Yield (name, child) pairs in a bottom-up manner
"""
for ch in self.children:
for el in ch.traverse():
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
def __iter__(
self,
loc: Optional[Location] = None,
name: Optional[str] = None,
color: Optional[Color] = None,
) -> Iterator[Tuple[Shape, str, Location, Optional[Color]]]:
"""
Assembly iterator yielding shapes, names, locations and colors.
"""
name = f"{name}/{self.name}" if name else self.name
loc = loc * self.loc if loc else self.loc
color = self.color if self.color else color
if self.obj:
yield self.obj if isinstance(self.obj, Shape) else Compound.makeCompound(
s for s in self.obj.vals() if isinstance(s, Shape)
), name, loc, color
for ch in self.children:
yield from ch.__iter__(loc, name, color)
def toCompound(self) -> Compound:
"""
Returns a Compound made from this Assembly (including all children) with the
current Locations applied. Usually this method would only be used after solving.
"""
shapes = self.shapes
shapes.extend((child.toCompound() for child in self.children))
return Compound.makeCompound(shapes).locate(self.loc)
def _repr_javascript_(self):
"""
Jupyter 3D representation support
"""
from .occ_impl.jupyter_tools import display
return display(self)._repr_javascript_()