Spaces:
Running
Running
| """Placement-based assembly system on top of CadQuery's Assembly. | |
| An assembly is a named collection of components; each component points to a | |
| part from the main scene plus a (location, rotation) placement. Solve() | |
| builds the cq.Assembly, optionally applies stub mate constraints, then writes | |
| the resulting compound back into the scene as a single part named | |
| '_asm_<asm>' so it shows up in the 3D viewer. | |
| """ | |
| from __future__ import annotations | |
| import cadquery as cq | |
| # tail-spec -> CadQuery tag-style face selector for box-like solids | |
| _FACE_TAG = { | |
| "top": ">Z", | |
| "bottom": "<Z", | |
| "left": "<X", | |
| "right": ">X", | |
| "front": "<Y", | |
| "back": ">Y", | |
| } | |
| def _resolve_selector(comp_obj: cq.Workplane, selector_tail: str): | |
| """Return either a CadQuery face-tag string or a concrete sub-shape. | |
| selector_tail is the part after '<comp>.', e.g. 'face.top', 'edge.3', | |
| 'axis.Z', 'point.0'. | |
| """ | |
| parts = selector_tail.split(".") | |
| if len(parts) != 2: | |
| raise ValueError(f"bad selector tail '{selector_tail}'") | |
| kind, spec = parts[0].lower(), parts[1] | |
| if kind == "face": | |
| s = spec.lower() | |
| if s not in _FACE_TAG: | |
| raise ValueError(f"face spec must be one of {list(_FACE_TAG)}, got '{spec}'") | |
| return _FACE_TAG[s] | |
| if kind == "edge": | |
| edges = comp_obj.edges().vals() | |
| return edges[int(spec)] | |
| if kind == "axis": | |
| ax = spec.upper() | |
| if ax not in ("X", "Y", "Z"): | |
| raise ValueError(f"axis must be X/Y/Z, got '{spec}'") | |
| bb = comp_obj.val().BoundingBox() | |
| center = cq.Vector((bb.xmin + bb.xmax) / 2, | |
| (bb.ymin + bb.ymax) / 2, | |
| (bb.zmin + bb.zmax) / 2) | |
| direction = {"X": cq.Vector(1, 0, 0), | |
| "Y": cq.Vector(0, 1, 0), | |
| "Z": cq.Vector(0, 0, 1)}[ax] | |
| return cq.Edge.makeLine(center, center + direction) | |
| if kind == "point": | |
| verts = comp_obj.vertices().vals() | |
| return verts[int(spec)] | |
| raise ValueError(f"unknown selector kind '{kind}'") | |
| def _split_selector(sel: str) -> tuple[str, str]: | |
| """Split '<comp>.<kind>.<spec>' -> ('<comp>', '<kind>.<spec>').""" | |
| first, _, rest = sel.partition(".") | |
| if not rest: | |
| raise ValueError(f"selector '{sel}' missing tail") | |
| return first, rest | |
| def _tag_for_selector(comp_obj: cq.Workplane, comp_name: str, sel_tail: str) -> str: | |
| """Build a CadQuery assembly tag string like 'name@faces@>Z' or 'name?tag'. | |
| Falls back to indexed face/edge tags via an OCCT-style query when needed. | |
| """ | |
| parts = sel_tail.split(".") | |
| kind, spec = parts[0].lower(), parts[1] | |
| if kind == "face": | |
| return f"{comp_name}@faces@{_FACE_TAG[spec.lower()]}" | |
| if kind == "edge": | |
| # CadQuery tag strings don't support raw indices; this query form | |
| # picks the i-th edge in the +X-sorted list, which is good enough. | |
| return f"{comp_name}@edges@>>X[{int(spec)}]" | |
| if kind == "axis": | |
| ax = spec.upper() | |
| return f"{comp_name}@faces@>{ax}" | |
| if kind == "point": | |
| return f"{comp_name}@vertices@>>X[{int(spec)}]" | |
| raise ValueError(f"unknown selector kind '{kind}'") | |
| class AssemblyEngine: | |
| def __init__(self, parts_ref: dict[str, cq.Workplane]): | |
| self.parts_ref = parts_ref | |
| # name -> { comps: {comp_name: {part, x, y, z, rx, ry, rz}}, mates: [...] } | |
| self.assemblies: dict[str, dict] = {} | |
| def _a(self, name: str) -> dict: | |
| if name not in self.assemblies: | |
| existing = list(self.assemblies) | |
| hint = f" → run 'asm new {name}' first to create it" | |
| if existing: | |
| hint += f" (existing assemblies: {existing})" | |
| raise KeyError(f"no assembly '{name}'.\n{hint}") | |
| return self.assemblies[name] | |
| # ----- structure ----- # | |
| def new(self, name: str) -> str: | |
| self.assemblies[name] = {"comps": {}, "mates": []} | |
| return f"created assembly '{name}'" | |
| def add_component(self, assembly: str, component: str, part: str, | |
| x: float = 0, y: float = 0, z: float = 0, | |
| rx: float = 0, ry: float = 0, rz: float = 0) -> str: | |
| # Friendly: auto-create the assembly if the user skipped 'asm new' | |
| if assembly not in self.assemblies: | |
| self.new(assembly) | |
| a = self.assemblies[assembly] | |
| if part not in self.parts_ref: | |
| available = list(self.parts_ref.keys()) | |
| hint = (f" → no part named '{part}' exists in the scene yet.\n" | |
| f" Build it first, e.g. 'box {part} 20 20 20' or " | |
| f"'cyl {part} 5 30'.") | |
| if available: | |
| hint += f"\n Existing parts: {available}" | |
| raise KeyError(hint) | |
| a["comps"][component] = { | |
| "part": part, | |
| "x": float(x), "y": float(y), "z": float(z), | |
| "rx": float(rx), "ry": float(ry), "rz": float(rz), | |
| } | |
| return (f"added component '{component}' (part '{part}') to " | |
| f"'{assembly}' at ({x},{y},{z}) rot ({rx},{ry},{rz})") | |
| def move_component(self, assembly: str, component: str, | |
| dx: float = 0, dy: float = 0, dz: float = 0) -> str: | |
| a = self._a(assembly) | |
| if component not in a["comps"]: | |
| raise KeyError(f"no component '{component}' in '{assembly}'") | |
| c = a["comps"][component] | |
| c["x"] += float(dx); c["y"] += float(dy); c["z"] += float(dz) | |
| return f"moved '{component}' by ({dx},{dy},{dz})" | |
| def rotate_component(self, assembly: str, component: str, | |
| rx: float = 0, ry: float = 0, rz: float = 0) -> str: | |
| a = self._a(assembly) | |
| if component not in a["comps"]: | |
| raise KeyError(f"no component '{component}'") | |
| c = a["comps"][component] | |
| c["rx"] += float(rx); c["ry"] += float(ry); c["rz"] += float(rz) | |
| return f"rotated '{component}' by ({rx},{ry},{rz}) deg" | |
| def remove_component(self, assembly: str, component: str) -> str: | |
| a = self._a(assembly) | |
| if component not in a["comps"]: | |
| raise KeyError(f"no component '{component}'") | |
| del a["comps"][component] | |
| return f"removed '{component}' from '{assembly}'" | |
| # ----- mates ----- # | |
| def add_mate(self, assembly: str, kind: str, a_sel: str, b_sel: str) -> str: | |
| """Record a mate; solved by solve() via cq.Assembly.constrain/solve.""" | |
| a = self._a(assembly) | |
| if kind not in ("Plane", "Axis", "Point"): | |
| raise ValueError(f"unknown mate kind '{kind}' (allowed: Plane, Axis, Point)") | |
| # validate selector shape early | |
| _split_selector(a_sel); _split_selector(b_sel) | |
| a["mates"].append({"kind": kind, "a": a_sel, "b": b_sel}) | |
| return f"recorded mate {kind} between '{a_sel}' and '{b_sel}'" | |
| def _placement_loc(self, c: dict) -> cq.Location: | |
| loc = cq.Location(cq.Vector(c["x"], c["y"], c["z"]), | |
| cq.Vector(0, 0, 1), c["rz"]) | |
| loc = loc * cq.Location(cq.Vector(0, 0, 0), cq.Vector(1, 0, 0), c["rx"]) | |
| loc = loc * cq.Location(cq.Vector(0, 0, 0), cq.Vector(0, 1, 0), c["ry"]) | |
| return loc | |
| def _write_back_location(self, c: dict, loc: cq.Location) -> None: | |
| """Decompose a solved Location into (xyz, rx,ry,rz deg) and store.""" | |
| try: | |
| (x, y, z), (rx, ry, rz) = loc.toTuple() | |
| c["x"], c["y"], c["z"] = float(x), float(y), float(z) | |
| c["rx"], c["ry"], c["rz"] = float(rx), float(ry), float(rz) | |
| except Exception: | |
| # fall back to translation-only readback | |
| t = loc.wrapped.Transformation().TranslationPart() | |
| c["x"], c["y"], c["z"] = float(t.X()), float(t.Y()), float(t.Z()) | |
| # ----- build ----- # | |
| def place_parts(self, assembly: str) -> str: | |
| """Move each component's underlying part to its placed location. | |
| Unlike solve(), no compound is created — the SOURCE parts physically | |
| relocate so they're visible at the right spots. This is what users | |
| expect when they 'build' an assembly. | |
| """ | |
| a = self._a(assembly) | |
| moved = [] | |
| for comp_name, c in a["comps"].items(): | |
| part_name = c["part"] | |
| if part_name not in self.parts_ref: | |
| continue | |
| part = self.parts_ref[part_name] | |
| # rotation first (about origin), then translation | |
| if c["rx"] != 0: | |
| part = part.rotate((0, 0, 0), (1, 0, 0), c["rx"]) | |
| if c["ry"] != 0: | |
| part = part.rotate((0, 0, 0), (0, 1, 0), c["ry"]) | |
| if c["rz"] != 0: | |
| part = part.rotate((0, 0, 0), (0, 0, 1), c["rz"]) | |
| part = part.translate((c["x"], c["y"], c["z"])) | |
| self.parts_ref[part_name] = part | |
| moved.append(f"'{part_name}' -> ({c['x']:.1f},{c['y']:.1f},{c['z']:.1f})") | |
| return (f"assembly '{assembly}': placed {len(moved)} component(s).\n " | |
| + "\n ".join(moved)) | |
| def solve(self, assembly: str) -> str: | |
| a = self._a(assembly) | |
| warn = "" | |
| asm = cq.Assembly(name=assembly) | |
| for comp_name, c in a["comps"].items(): | |
| asm.add(self.parts_ref[c["part"]], name=comp_name, | |
| loc=self._placement_loc(c)) | |
| if a["mates"]: | |
| try: | |
| for m in a["mates"]: | |
| a_comp, a_tail = _split_selector(m["a"]) | |
| b_comp, b_tail = _split_selector(m["b"]) | |
| a_obj = self.parts_ref[a["comps"][a_comp]["part"]] | |
| b_obj = self.parts_ref[a["comps"][b_comp]["part"]] | |
| a_tag = _tag_for_selector(a_obj, a_comp, a_tail) | |
| b_tag = _tag_for_selector(b_obj, b_comp, b_tail) | |
| asm.constrain(a_tag, b_tag, m["kind"]) | |
| asm.solve() | |
| # read back solved locations | |
| for comp_name, c in a["comps"].items(): | |
| obj = asm.objects.get(comp_name) | |
| if obj is not None and obj.loc is not None: | |
| self._write_back_location(c, obj.loc) | |
| except Exception as e: | |
| warn = (f"WARNING: mate solve failed ({e}); " | |
| f"positioning by placement only. ") | |
| # rebuild asm without constraints so toCompound uses placements | |
| asm = cq.Assembly(name=assembly) | |
| for comp_name, c in a["comps"].items(): | |
| asm.add(self.parts_ref[c["part"]], name=comp_name, | |
| loc=self._placement_loc(c)) | |
| compound = asm.toCompound() | |
| self.parts_ref[f"_asm_{assembly}"] = cq.Workplane(obj=compound) | |
| return (f"{warn}built assembly '{assembly}' with {len(a['comps'])} " | |
| f"components -> part '_asm_{assembly}'") | |
| # ----- access ----- # | |
| def list_assemblies(self) -> str: | |
| if not self.assemblies: | |
| return "no assemblies" | |
| out = [] | |
| for n, a in self.assemblies.items(): | |
| out.append(f" {n}: {len(a['comps'])} components, {len(a['mates'])} mates") | |
| return "assemblies:\n" + "\n".join(out) | |
| def info(self, name: str) -> str: | |
| a = self._a(name) | |
| lines = [f"assembly '{name}'"] | |
| for cn, c in a["comps"].items(): | |
| lines.append(f" {cn}: part='{c['part']}' " | |
| f"loc=({c['x']:.2f},{c['y']:.2f},{c['z']:.2f}) " | |
| f"rot=({c['rx']:.1f},{c['ry']:.1f},{c['rz']:.1f})") | |
| for i, m in enumerate(a["mates"]): | |
| lines.append(f" mate{i}: {m['kind']} {m['a']} <-> {m['b']}") | |
| return "\n".join(lines) | |
| def delete(self, name: str) -> str: | |
| self._a(name) | |
| del self.assemblies[name] | |
| # drop any auto-built compound from the scene | |
| self.parts_ref.pop(f"_asm_{name}", None) | |
| return f"deleted assembly '{name}'" | |
| def export_step(self, assembly: str, output_dir: str, | |
| filename: str | None = None) -> str: | |
| import os | |
| a = self._a(assembly) | |
| asm = cq.Assembly(name=assembly) | |
| for comp_name, c in a["comps"].items(): | |
| part = self.parts_ref[c["part"]] | |
| loc = cq.Location(cq.Vector(c["x"], c["y"], c["z"]), | |
| cq.Vector(0, 0, 1), c["rz"]) | |
| loc = loc * cq.Location(cq.Vector(0, 0, 0), cq.Vector(1, 0, 0), c["rx"]) | |
| loc = loc * cq.Location(cq.Vector(0, 0, 0), cq.Vector(0, 1, 0), c["ry"]) | |
| asm.add(part, name=comp_name, loc=loc) | |
| fname = filename or f"{assembly}.step" | |
| path = os.path.join(output_dir, fname) | |
| asm.save(path) | |
| return path | |
| def dispatch_assembly(eng: AssemblyEngine, op: str, args: dict) -> str: | |
| fn = getattr(eng, op, None) | |
| if fn is None or op.startswith("_"): | |
| raise ValueError(f"unknown assembly op '{op}'") | |
| return fn(**args) | |