mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
284 lines
9.3 KiB
Python
284 lines
9.3 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Sequence
|
|
|
|
from common.assembly_flatten import CatalogEntry, filesystem_entry
|
|
from common.assembly_spec import (
|
|
REPO_ROOT,
|
|
AssemblySpec,
|
|
AssemblyNodeSpec,
|
|
assembly_spec_from_payload,
|
|
assembly_spec_children,
|
|
)
|
|
from common.catalog import find_source_by_cad_ref
|
|
from common.render import part_selector_manifest_path
|
|
|
|
|
|
GIT_LFS_POINTER_PREFIX = b"version https://git-lfs.github.com/spec/v1\n"
|
|
|
|
|
|
def _relative_to_repo(path: Path) -> str:
|
|
resolved = path.resolve()
|
|
try:
|
|
return resolved.relative_to(REPO_ROOT).as_posix()
|
|
except ValueError:
|
|
return resolved.as_posix()
|
|
|
|
|
|
def _is_git_lfs_pointer(path: Path) -> bool:
|
|
try:
|
|
with path.open("rb") as handle:
|
|
return handle.read(len(GIT_LFS_POINTER_PREFIX)) == GIT_LFS_POINTER_PREFIX
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def _location_from_transform(transform: tuple[float, ...]):
|
|
import build123d
|
|
from OCP.gp import gp_Trsf
|
|
|
|
trsf = gp_Trsf()
|
|
trsf.SetValues(
|
|
transform[0],
|
|
transform[1],
|
|
transform[2],
|
|
transform[3],
|
|
transform[4],
|
|
transform[5],
|
|
transform[6],
|
|
transform[7],
|
|
transform[8],
|
|
transform[9],
|
|
transform[10],
|
|
transform[11],
|
|
)
|
|
return build123d.Location(trsf)
|
|
|
|
|
|
def _component_name(instance_path: tuple[str, ...]) -> str:
|
|
return "__".join(instance_path) or "root"
|
|
|
|
|
|
def _load_step_shape(step_path: Path):
|
|
if not step_path.exists():
|
|
raise FileNotFoundError(f"Referenced STEP file is missing: {_relative_to_repo(step_path)}")
|
|
if _is_git_lfs_pointer(step_path):
|
|
raise RuntimeError(f"Referenced STEP file is a Git LFS pointer: {_relative_to_repo(step_path)}")
|
|
|
|
import build123d
|
|
|
|
try:
|
|
return build123d.import_step(step_path)
|
|
except Exception as exc:
|
|
raise RuntimeError(f"Failed to load referenced STEP file: {_relative_to_repo(step_path)}") from exc
|
|
|
|
|
|
def _step_has_assembly_artifact(step_path: Path) -> bool:
|
|
import json
|
|
|
|
topology_path = part_selector_manifest_path(step_path)
|
|
try:
|
|
payload = json.loads(topology_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError):
|
|
return False
|
|
root = payload.get("assembly", {}).get("root") if isinstance(payload.get("assembly"), dict) else None
|
|
return isinstance(root, dict) and bool(root.get("children"))
|
|
|
|
|
|
def _load_step_assembly_shape(step_path: Path, *, label: str):
|
|
import build123d
|
|
|
|
from common.step_scene import load_step_scene, occurrence_selector_id, scene_occurrence_shape
|
|
|
|
scene = load_step_scene(step_path)
|
|
|
|
def node_label(node: object) -> str:
|
|
return str(getattr(node, "name", None) or getattr(node, "source_name", None) or occurrence_selector_id(node)).strip()
|
|
|
|
def build_node(node: object):
|
|
children = list(getattr(node, "children", []) or [])
|
|
if children:
|
|
child_shapes = [build_node(child) for child in children]
|
|
return build123d.Compound(
|
|
obj=child_shapes,
|
|
children=child_shapes,
|
|
label=node_label(node),
|
|
)
|
|
shape = build123d.Shape(obj=scene_occurrence_shape(scene, node))
|
|
shape.label = node_label(node)
|
|
return shape
|
|
|
|
roots = [build_node(root) for root in scene.roots]
|
|
if not roots:
|
|
return _load_step_shape(step_path)
|
|
return build123d.Compound(obj=roots, children=roots, label=label)
|
|
|
|
|
|
def _source_color_for_cad_ref(cad_ref: str):
|
|
try:
|
|
source = find_source_by_cad_ref(cad_ref)
|
|
except Exception:
|
|
source = None
|
|
return source.color if source is not None else None
|
|
|
|
|
|
def _clear_shape_colors(shape: object) -> None:
|
|
if hasattr(shape, "color"):
|
|
shape.color = None
|
|
for child in getattr(shape, "children", []) or []:
|
|
_clear_shape_colors(child)
|
|
|
|
|
|
def _apply_source_color(shape: object, cad_ref: str, *, use_source_colors: bool) -> None:
|
|
import build123d
|
|
|
|
source_color = _source_color_for_cad_ref(cad_ref)
|
|
if not use_source_colors:
|
|
_clear_shape_colors(shape)
|
|
elif source_color is not None:
|
|
shape.color = build123d.Color(*source_color)
|
|
|
|
|
|
def _shape_for_part_entry(entry: CatalogEntry, *, label: str, use_source_colors: bool):
|
|
step_path = entry.step_path.resolve() if entry.step_path is not None else None
|
|
if step_path is None:
|
|
raise RuntimeError(f"Part source {entry.source_ref} is missing STEP source path")
|
|
shape = (
|
|
_load_step_assembly_shape(step_path, label=label)
|
|
if _step_has_assembly_artifact(step_path)
|
|
else _load_step_shape(step_path)
|
|
)
|
|
_apply_source_color(shape, entry.cad_ref, use_source_colors=use_source_colors)
|
|
shape.label = label
|
|
return shape
|
|
|
|
|
|
def _build_node_shape(
|
|
node: AssemblyNodeSpec,
|
|
*,
|
|
resolve_entry,
|
|
instance_path: tuple[str, ...],
|
|
parent_use_source_colors: bool,
|
|
stack: tuple[str, ...],
|
|
):
|
|
import build123d
|
|
|
|
component_path = (*instance_path, node.instance_id) if instance_path else (node.instance_id,)
|
|
label = _component_name(component_path)
|
|
use_source_colors = parent_use_source_colors and node.use_source_colors
|
|
|
|
if node.children:
|
|
child_shapes = [
|
|
_build_node_shape(
|
|
child,
|
|
resolve_entry=resolve_entry,
|
|
instance_path=component_path,
|
|
parent_use_source_colors=use_source_colors,
|
|
stack=stack,
|
|
)
|
|
for child in node.children
|
|
]
|
|
shape = build123d.Compound(obj=child_shapes, children=child_shapes, label=label)
|
|
else:
|
|
if node.source_path is None or node.path is None:
|
|
raise RuntimeError(f"Assembly node {label} is missing a STEP source path")
|
|
child_entry = resolve_entry(node.source_path)
|
|
if child_entry is None:
|
|
raise RuntimeError(f"Assembly node {label} references missing CAD source {node.path}")
|
|
if child_entry.kind == "assembly":
|
|
if child_entry.assembly_spec is None:
|
|
raise RuntimeError(f"Assembly source {child_entry.source_ref} is missing assembly spec data")
|
|
stack_key = child_entry.source_ref
|
|
if stack_key in stack:
|
|
cycle = " -> ".join((*stack, stack_key))
|
|
raise RuntimeError(f"Assembly cycle detected: {cycle}")
|
|
shape = _compound_from_nodes(
|
|
assembly_spec_children(child_entry.assembly_spec),
|
|
label=label,
|
|
resolve_entry=resolve_entry,
|
|
instance_path=component_path,
|
|
parent_use_source_colors=use_source_colors,
|
|
stack=(*stack, stack_key),
|
|
)
|
|
elif child_entry.kind == "part":
|
|
shape = _shape_for_part_entry(child_entry, label=label, use_source_colors=use_source_colors)
|
|
else:
|
|
raise RuntimeError(f"Assembly node {label} resolved to unsupported CAD source kind: {child_entry.kind}")
|
|
|
|
moved = shape.moved(_location_from_transform(node.transform))
|
|
moved.label = label
|
|
if not use_source_colors:
|
|
_clear_shape_colors(moved)
|
|
return moved
|
|
|
|
|
|
def _compound_from_nodes(
|
|
nodes: Sequence[AssemblyNodeSpec],
|
|
*,
|
|
label: str,
|
|
resolve_entry,
|
|
instance_path: tuple[str, ...] = (),
|
|
parent_use_source_colors: bool,
|
|
stack: tuple[str, ...],
|
|
):
|
|
import build123d
|
|
|
|
children = [
|
|
_build_node_shape(
|
|
node,
|
|
resolve_entry=resolve_entry,
|
|
instance_path=instance_path,
|
|
parent_use_source_colors=parent_use_source_colors,
|
|
stack=stack,
|
|
)
|
|
for node in nodes
|
|
]
|
|
if not children:
|
|
raise RuntimeError(f"Assembly {label} has no resolved STEP instances")
|
|
return build123d.Compound(
|
|
obj=children,
|
|
children=children,
|
|
label=label,
|
|
)
|
|
|
|
|
|
def build_assembly_compound(assembly_spec: AssemblySpec, *, label: str | None = None):
|
|
root_source = filesystem_entry(assembly_spec.assembly_path)
|
|
root_source_ref = root_source.source_ref if root_source is not None else _relative_to_repo(assembly_spec.assembly_path)
|
|
return _compound_from_nodes(
|
|
assembly_spec_children(assembly_spec),
|
|
label=label or Path(assembly_spec.assembly_path).stem,
|
|
resolve_entry=filesystem_entry,
|
|
parent_use_source_colors=True,
|
|
stack=(root_source_ref,),
|
|
)
|
|
|
|
|
|
def export_assembly_step(assembly_spec: AssemblySpec, output_path: Path) -> Path:
|
|
import build123d
|
|
|
|
assembly = build_assembly_compound(assembly_spec, label=output_path.stem)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
try:
|
|
success = build123d.export_step(assembly, output_path)
|
|
except Exception as exc:
|
|
raise RuntimeError(f"Failed to write assembly STEP file: {_relative_to_repo(output_path)}") from exc
|
|
if not success:
|
|
raise RuntimeError(f"Failed to write assembly STEP file: {_relative_to_repo(output_path)}")
|
|
if not output_path.exists() or output_path.stat().st_size <= 0:
|
|
raise RuntimeError(f"Assembly STEP export did not create {_relative_to_repo(output_path)}")
|
|
print(f"Wrote STEP: {output_path}")
|
|
return output_path
|
|
|
|
|
|
def export_assembly_step_from_payload(
|
|
payload: object,
|
|
*,
|
|
assembly_path: Path,
|
|
output_path: Path,
|
|
) -> Path:
|
|
return export_assembly_step(
|
|
assembly_spec_from_payload(assembly_path, payload),
|
|
output_path,
|
|
)
|