mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
600 lines
21 KiB
Python
600 lines
21 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from dataclasses import dataclass
|
|
from fnmatch import fnmatch
|
|
from pathlib import Path, PurePosixPath
|
|
|
|
from .metadata import GeneratorMetadata, normalize_mesh_numeric, parse_generator_metadata
|
|
|
|
|
|
REPO_ROOT = Path.cwd().resolve()
|
|
CAD_ROOT = REPO_ROOT
|
|
STEP_SUFFIXES = (".step", ".stp")
|
|
VIEWER_ARTIFACT_FILENAMES = {
|
|
".glb": "model.glb",
|
|
".topology.json": "topology.json",
|
|
".topology.bin": "topology.bin",
|
|
}
|
|
IGNORED_DISCOVERY_DIR_NAMES = {
|
|
"__pycache__",
|
|
".cache",
|
|
".eggs",
|
|
".env",
|
|
".git",
|
|
".hg",
|
|
".mypy_cache",
|
|
".pytest_cache",
|
|
".ruff_cache",
|
|
".svn",
|
|
".tox",
|
|
".venv",
|
|
"build",
|
|
"dist",
|
|
"env",
|
|
"node_modules",
|
|
"site-packages",
|
|
"venv",
|
|
}
|
|
GENERATOR_NAME_MARKERS = (b"gen_step", b"gen_dxf", b"gen_urdf")
|
|
|
|
|
|
class CadSourceError(ValueError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StepImportOptions:
|
|
export_stl: bool = False
|
|
stl_output: str | None = None
|
|
stl_tolerance: float | None = None
|
|
stl_angular_tolerance: float | None = None
|
|
glb_tolerance: float | None = None
|
|
glb_angular_tolerance: float | None = None
|
|
color: tuple[float, float, float, float] | None = None
|
|
skip_topology: bool = False
|
|
|
|
@property
|
|
def has_metadata(self) -> bool:
|
|
return any(
|
|
(
|
|
self.export_stl,
|
|
self.stl_output is not None,
|
|
self.stl_tolerance is not None,
|
|
self.stl_angular_tolerance is not None,
|
|
self.glb_tolerance is not None,
|
|
self.glb_angular_tolerance is not None,
|
|
self.color is not None,
|
|
self.skip_topology,
|
|
)
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CadSource:
|
|
source_ref: str
|
|
cad_ref: str
|
|
kind: str
|
|
source_path: Path
|
|
source: str
|
|
origin_path: Path
|
|
script_path: Path | None = None
|
|
generator_metadata: GeneratorMetadata | None = None
|
|
step_path: Path | None = None
|
|
stl_path: Path | None = None
|
|
dxf_path: Path | None = None
|
|
urdf_path: Path | None = None
|
|
export_stl: bool = False
|
|
stl_tolerance: float | None = None
|
|
stl_angular_tolerance: float | None = None
|
|
glb_tolerance: float | None = None
|
|
glb_angular_tolerance: float | None = None
|
|
color: tuple[float, float, float, float] | None = None
|
|
skip_topology: bool = False
|
|
|
|
@property
|
|
def selector_manifest_path(self) -> Path | None:
|
|
return (
|
|
viewer_artifact_path_for_step_path(self.step_path, ".topology.json")
|
|
if self.step_path is not None and not self.skip_topology
|
|
else None
|
|
)
|
|
|
|
@property
|
|
def selector_binary_path(self) -> Path | None:
|
|
return (
|
|
viewer_artifact_path_for_step_path(self.step_path, ".topology.bin")
|
|
if self.step_path is not None and not self.skip_topology
|
|
else None
|
|
)
|
|
|
|
@property
|
|
def glb_path(self) -> Path | None:
|
|
return viewer_artifact_path_for_step_path(self.step_path, ".glb") if self.step_path is not None else None
|
|
|
|
@property
|
|
def generated_paths(self) -> tuple[Path, ...]:
|
|
paths: list[Path] = []
|
|
if self.source == "generated":
|
|
if self.step_path is not None:
|
|
paths.append(self.step_path)
|
|
if self.dxf_path is not None:
|
|
paths.append(self.dxf_path)
|
|
if self.urdf_path is not None:
|
|
paths.append(self.urdf_path)
|
|
if self.export_stl and self.stl_path is not None:
|
|
paths.append(self.stl_path)
|
|
if self.glb_path is not None:
|
|
paths.append(self.glb_path)
|
|
if self.selector_manifest_path is not None:
|
|
paths.append(self.selector_manifest_path)
|
|
if self.selector_binary_path is not None:
|
|
paths.append(self.selector_binary_path)
|
|
return tuple(path.resolve() for path in paths)
|
|
|
|
|
|
def iter_cad_sources(root: Path | None = None) -> tuple[CadSource, ...]:
|
|
root = CAD_ROOT if root is None else root
|
|
resolved_root = root.resolve()
|
|
python_sources = _iter_python_sources(resolved_root)
|
|
generated_step_paths = {
|
|
source.step_path.resolve()
|
|
for source in python_sources
|
|
if source.step_path is not None
|
|
}
|
|
sources = [
|
|
*python_sources,
|
|
*_iter_step_sources(resolved_root, excluded_step_paths=generated_step_paths),
|
|
]
|
|
by_cad_ref: dict[str, CadSource] = {}
|
|
by_source_ref: dict[str, CadSource] = {}
|
|
by_step_path: dict[Path, CadSource] = {}
|
|
by_generated_path: dict[Path, CadSource] = {}
|
|
for source in sources:
|
|
existing = by_cad_ref.get(source.cad_ref)
|
|
if existing is not None:
|
|
raise CadSourceError(
|
|
"Duplicate CAD STEP ref "
|
|
f"{source.cad_ref!r}: {_source_label(existing)} and {_source_label(source)}"
|
|
)
|
|
by_cad_ref[source.cad_ref] = source
|
|
existing_source = by_source_ref.get(source.source_ref)
|
|
if existing_source is not None:
|
|
raise CadSourceError(
|
|
"Duplicate CAD source ref "
|
|
f"{source.source_ref!r}: {_source_label(existing_source)} and {_source_label(source)}"
|
|
)
|
|
by_source_ref[source.source_ref] = source
|
|
if source.step_path is not None:
|
|
existing_step = by_step_path.get(source.step_path.resolve())
|
|
if existing_step is not None:
|
|
raise CadSourceError(
|
|
"Duplicate CAD STEP source "
|
|
f"{_relative_to_repo(source.step_path)}: {_source_label(existing_step)} and {_source_label(source)}"
|
|
)
|
|
by_step_path[source.step_path.resolve()] = source
|
|
for generated_path in source.generated_paths:
|
|
resolved_generated_path = generated_path.resolve()
|
|
existing_generated = by_generated_path.get(resolved_generated_path)
|
|
if existing_generated is not None and existing_generated.source_ref != source.source_ref:
|
|
raise CadSourceError(
|
|
"Duplicate CAD generated output "
|
|
f"{_relative_to_repo(generated_path)}: "
|
|
f"{_source_label(existing_generated)} and {_source_label(source)}"
|
|
)
|
|
by_generated_path[resolved_generated_path] = source
|
|
return tuple(sorted(by_cad_ref.values(), key=lambda source: source.source_ref))
|
|
|
|
|
|
def source_from_path(
|
|
path: Path,
|
|
*,
|
|
step_kind: str = "part",
|
|
step_options: StepImportOptions | None = None,
|
|
) -> CadSource | None:
|
|
resolved = path.resolve()
|
|
if resolved.suffix.lower() == ".py":
|
|
return _read_python_source(resolved)
|
|
if resolved.suffix.lower() in STEP_SUFFIXES:
|
|
return _read_step_source(resolved, kind=step_kind, options=step_options)
|
|
return None
|
|
|
|
|
|
def source_by_cad_ref(root: Path | None = None) -> dict[str, CadSource]:
|
|
return {source.cad_ref: source for source in iter_cad_sources(root)}
|
|
|
|
|
|
def find_source_by_cad_ref(cad_ref: str, root: Path | None = None) -> CadSource | None:
|
|
normalized = normalize_cad_ref(cad_ref)
|
|
return source_by_cad_ref(root).get(normalized or "")
|
|
|
|
|
|
def find_source_by_source_ref(source_ref: str, root: Path | None = None) -> CadSource | None:
|
|
normalized = normalize_source_ref(source_ref)
|
|
if not normalized:
|
|
return None
|
|
for source in iter_cad_sources(root):
|
|
if source.source_ref == normalized:
|
|
return source
|
|
return None
|
|
|
|
|
|
def find_source_by_path(path: Path, root: Path | None = None) -> CadSource | None:
|
|
resolved_path = path.resolve()
|
|
for source in iter_cad_sources(root):
|
|
paths = [
|
|
source.source_path,
|
|
source.step_path,
|
|
source.script_path,
|
|
source.dxf_path,
|
|
source.urdf_path,
|
|
*source.generated_paths,
|
|
]
|
|
if any(candidate is not None and candidate.resolve() == resolved_path for candidate in paths):
|
|
return source
|
|
return None
|
|
|
|
|
|
def source_ref_from_path(path: Path) -> str:
|
|
resolved = path.resolve()
|
|
try:
|
|
relative = resolved.relative_to(CAD_ROOT.resolve())
|
|
except ValueError:
|
|
return resolved.as_posix()
|
|
return relative.as_posix()
|
|
|
|
|
|
def cad_ref_from_step_path(path: Path) -> str:
|
|
resolved = path.resolve()
|
|
try:
|
|
relative = resolved.relative_to(CAD_ROOT.resolve())
|
|
except ValueError:
|
|
relative = PurePosixPath(resolved.as_posix())
|
|
name = relative.name
|
|
suffix = relative.suffix.lower()
|
|
if suffix in STEP_SUFFIXES:
|
|
return relative.with_suffix("").as_posix()
|
|
raise CadSourceError(f"{_relative_to_repo(path)} is not a CAD STEP source")
|
|
|
|
|
|
def normalize_source_ref(raw_ref: str) -> str | None:
|
|
normalized = str(raw_ref or "").replace("\\", "/").strip().strip("/")
|
|
if not normalized:
|
|
return None
|
|
parts = normalized.split("/")
|
|
if any(not part or part in {".", ".."} for part in parts):
|
|
return None
|
|
return "/".join(parts)
|
|
|
|
|
|
def normalize_cad_ref(raw_ref: str) -> str | None:
|
|
normalized = normalize_source_ref(raw_ref)
|
|
if not normalized:
|
|
return None
|
|
suffix = PurePosixPath(normalized).suffix.lower()
|
|
if suffix in {".py", *STEP_SUFFIXES}:
|
|
normalized = str(PurePosixPath(normalized).with_suffix(""))
|
|
return normalized
|
|
|
|
|
|
def artifact_path_for_step_path(step_path: Path, suffix: str) -> Path:
|
|
return step_path.resolve().with_suffix(suffix)
|
|
|
|
|
|
def hidden_artifact_path_for_step_path(step_path: Path, suffix: str) -> Path:
|
|
base = step_path.resolve()
|
|
return base.with_name(f".{base.stem}{suffix}").resolve()
|
|
|
|
|
|
def viewer_directory_for_step_path(step_path: Path) -> Path:
|
|
base = step_path.resolve()
|
|
return (base.parent / f".{base.name}").resolve()
|
|
|
|
|
|
def viewer_artifact_path_for_step_path(step_path: Path, suffix: str) -> Path:
|
|
base = step_path.resolve()
|
|
artifact_name = VIEWER_ARTIFACT_FILENAMES.get(suffix)
|
|
if artifact_name is None:
|
|
raise ValueError(f"Unsupported STEP viewer artifact suffix: {suffix}")
|
|
return (viewer_directory_for_step_path(base) / artifact_name).resolve()
|
|
|
|
|
|
def _iter_python_sources(root: Path) -> tuple[CadSource, ...]:
|
|
sources: list[CadSource] = []
|
|
for script_path in _iter_paths(root, "*.py"):
|
|
if not _looks_like_generator_script(script_path):
|
|
continue
|
|
source = _read_python_source(script_path)
|
|
if source is not None:
|
|
sources.append(source)
|
|
return tuple(sources)
|
|
|
|
|
|
def _read_python_source(script_path: Path) -> CadSource | None:
|
|
resolved_script_path = script_path.resolve()
|
|
metadata = parse_generator_metadata(resolved_script_path)
|
|
if metadata is None:
|
|
return None
|
|
if metadata.kind not in {"part", "assembly"}:
|
|
raise CadSourceError(
|
|
f"{_relative_to_repo(resolved_script_path)} must define a part or assembly gen_step() entry"
|
|
)
|
|
if metadata.kind == "assembly" and metadata.skip_topology:
|
|
raise CadSourceError(
|
|
f"{_relative_to_repo(resolved_script_path)} skip_topology is not supported for assembly entries"
|
|
)
|
|
if metadata.step_output is None:
|
|
raise CadSourceError(
|
|
f"{_relative_to_repo(resolved_script_path)} step_output is required"
|
|
)
|
|
step_path = _resolve_configured_artifact_path(
|
|
metadata.step_output,
|
|
base_path=resolved_script_path,
|
|
default_path=None,
|
|
expected_suffixes=(".step",),
|
|
field_name="step_output",
|
|
)
|
|
if metadata.stl_output is not None and not metadata.export_stl:
|
|
raise CadSourceError(
|
|
f"{_relative_to_repo(resolved_script_path)} stl_output requires export_stl = True"
|
|
)
|
|
dxf_path = (
|
|
_resolve_configured_artifact_path(
|
|
_required_output(metadata.dxf_output, script_path=resolved_script_path, field_name="dxf_output"),
|
|
base_path=resolved_script_path,
|
|
default_path=None,
|
|
expected_suffixes=(".dxf",),
|
|
field_name="dxf_output",
|
|
)
|
|
if metadata.has_gen_dxf
|
|
else None
|
|
)
|
|
urdf_path = (
|
|
_resolve_configured_artifact_path(
|
|
_required_output(metadata.urdf_output, script_path=resolved_script_path, field_name="urdf_output"),
|
|
base_path=resolved_script_path,
|
|
default_path=None,
|
|
expected_suffixes=(".urdf",),
|
|
field_name="urdf_output",
|
|
)
|
|
if metadata.has_gen_urdf
|
|
else None
|
|
)
|
|
stl_path = (
|
|
_resolve_configured_artifact_path(
|
|
_required_output(metadata.stl_output, script_path=resolved_script_path, field_name="stl_output"),
|
|
base_path=resolved_script_path,
|
|
default_path=None,
|
|
expected_suffixes=(".stl",),
|
|
field_name="stl_output",
|
|
)
|
|
if metadata.export_stl
|
|
else None
|
|
)
|
|
return CadSource(
|
|
source_ref=source_ref_from_path(resolved_script_path),
|
|
cad_ref=cad_ref_from_step_path(step_path),
|
|
kind=metadata.kind,
|
|
source_path=resolved_script_path,
|
|
source="generated",
|
|
origin_path=resolved_script_path,
|
|
script_path=resolved_script_path,
|
|
generator_metadata=metadata,
|
|
step_path=step_path,
|
|
stl_path=stl_path,
|
|
dxf_path=dxf_path,
|
|
urdf_path=urdf_path,
|
|
export_stl=metadata.export_stl,
|
|
stl_tolerance=metadata.stl_tolerance,
|
|
stl_angular_tolerance=metadata.stl_angular_tolerance,
|
|
glb_tolerance=metadata.glb_tolerance,
|
|
glb_angular_tolerance=metadata.glb_angular_tolerance,
|
|
skip_topology=metadata.skip_topology,
|
|
)
|
|
|
|
|
|
def _iter_step_sources(root: Path, *, excluded_step_paths: set[Path]) -> tuple[CadSource, ...]:
|
|
sources: list[CadSource] = []
|
|
for pattern in ("*.step", "*.stp"):
|
|
for step_path in _iter_paths(root, pattern):
|
|
if step_path.resolve() in excluded_step_paths:
|
|
continue
|
|
sources.append(_read_step_source(step_path, kind="part"))
|
|
return tuple(sorted(sources, key=lambda source: source.source_ref))
|
|
|
|
|
|
def _read_step_source(
|
|
step_path: Path,
|
|
*,
|
|
kind: str,
|
|
options: StepImportOptions | None = None,
|
|
) -> CadSource:
|
|
resolved_step_path = step_path.resolve()
|
|
options = options or StepImportOptions()
|
|
if kind not in {"part", "assembly"}:
|
|
raise CadSourceError(f"{_relative_to_repo(resolved_step_path)} kind must be 'part' or 'assembly'")
|
|
if kind == "assembly" and options.skip_topology:
|
|
raise CadSourceError(
|
|
f"{_relative_to_repo(resolved_step_path)} skip_topology is not supported for assembly entries"
|
|
)
|
|
if resolved_step_path.suffix.lower() not in STEP_SUFFIXES:
|
|
raise CadSourceError(f"{_relative_to_repo(resolved_step_path)} source must end in .step or .stp")
|
|
if not resolved_step_path.is_file():
|
|
raise CadSourceError(
|
|
f"{_relative_to_repo(resolved_step_path)} source does not exist"
|
|
)
|
|
if options.stl_output is not None and not options.export_stl:
|
|
raise CadSourceError(
|
|
f"{_relative_to_repo(resolved_step_path)} stl_output requires export_stl = true"
|
|
)
|
|
if options.export_stl and options.stl_output is None:
|
|
raise CadSourceError(
|
|
f"{_relative_to_repo(resolved_step_path)} stl_output is required when export_stl = true"
|
|
)
|
|
stl_path = (
|
|
_resolve_configured_artifact_path(
|
|
options.stl_output,
|
|
base_path=resolved_step_path,
|
|
default_path=None,
|
|
expected_suffixes=(".stl",),
|
|
field_name="stl_output",
|
|
)
|
|
if options.export_stl
|
|
else None
|
|
)
|
|
|
|
cad_ref = cad_ref_from_step_path(resolved_step_path)
|
|
|
|
return CadSource(
|
|
source_ref=source_ref_from_path(resolved_step_path),
|
|
cad_ref=cad_ref,
|
|
kind=str(kind),
|
|
source_path=resolved_step_path,
|
|
source="imported",
|
|
origin_path=resolved_step_path,
|
|
step_path=resolved_step_path,
|
|
stl_path=stl_path,
|
|
export_stl=options.export_stl,
|
|
stl_tolerance=normalize_step_numeric(
|
|
options.stl_tolerance,
|
|
base_path=resolved_step_path,
|
|
field_name="stl_tolerance",
|
|
),
|
|
stl_angular_tolerance=normalize_step_numeric(
|
|
options.stl_angular_tolerance,
|
|
base_path=resolved_step_path,
|
|
field_name="stl_angular_tolerance",
|
|
),
|
|
glb_tolerance=normalize_step_numeric(
|
|
options.glb_tolerance,
|
|
base_path=resolved_step_path,
|
|
field_name="glb_tolerance",
|
|
),
|
|
glb_angular_tolerance=normalize_step_numeric(
|
|
options.glb_angular_tolerance,
|
|
base_path=resolved_step_path,
|
|
field_name="glb_angular_tolerance",
|
|
),
|
|
color=options.color,
|
|
skip_topology=options.skip_topology,
|
|
)
|
|
|
|
|
|
def _iter_paths(root: Path, pattern: str) -> tuple[Path, ...]:
|
|
paths: list[Path] = []
|
|
for current_root, dirnames, filenames in os.walk(root):
|
|
dirnames[:] = sorted(
|
|
dirname
|
|
for dirname in dirnames
|
|
if dirname not in IGNORED_DISCOVERY_DIR_NAMES
|
|
)
|
|
for filename in sorted(filenames):
|
|
if not fnmatch(filename, pattern):
|
|
continue
|
|
path = (Path(current_root) / filename).resolve()
|
|
if path.is_file():
|
|
paths.append(path)
|
|
return tuple(paths)
|
|
|
|
|
|
def _looks_like_generator_script(script_path: Path) -> bool:
|
|
try:
|
|
source_bytes = script_path.read_bytes()
|
|
except OSError:
|
|
return False
|
|
return any(marker in source_bytes for marker in GENERATOR_NAME_MARKERS)
|
|
|
|
|
|
def normalize_step_numeric(raw_value: object, *, base_path: Path, field_name: str) -> float | None:
|
|
try:
|
|
return normalize_mesh_numeric(raw_value, field_name=field_name)
|
|
except ValueError as exc:
|
|
raise CadSourceError(f"{_relative_to_repo(base_path)} {exc}") from exc
|
|
|
|
|
|
def normalize_step_color(
|
|
raw_value: object,
|
|
*,
|
|
base_path: Path,
|
|
field_name: str,
|
|
) -> tuple[float, float, float, float] | None:
|
|
if raw_value is None:
|
|
return None
|
|
if isinstance(raw_value, str):
|
|
value = raw_value.strip()
|
|
if value.startswith("#"):
|
|
value = value[1:]
|
|
if len(value) not in {6, 8}:
|
|
raise CadSourceError(f"{_relative_to_repo(base_path)} {field_name} must be #RRGGBB or #RRGGBBAA")
|
|
try:
|
|
components = [int(value[index : index + 2], 16) / 255.0 for index in range(0, len(value), 2)]
|
|
except ValueError as exc:
|
|
raise CadSourceError(f"{_relative_to_repo(base_path)} {field_name} must be valid hex") from exc
|
|
elif isinstance(raw_value, list) and len(raw_value) in {3, 4}:
|
|
components = []
|
|
for component in raw_value:
|
|
try:
|
|
number = float(component)
|
|
except (TypeError, ValueError) as exc:
|
|
raise CadSourceError(
|
|
f"{_relative_to_repo(base_path)} {field_name} components must be numeric"
|
|
) from exc
|
|
if not 0.0 <= number <= 1.0:
|
|
raise CadSourceError(
|
|
f"{_relative_to_repo(base_path)} {field_name} components must be between 0 and 1"
|
|
)
|
|
components.append(number)
|
|
else:
|
|
raise CadSourceError(f"{_relative_to_repo(base_path)} {field_name} must be an RGB/RGBA array or hex string")
|
|
if len(components) == 3:
|
|
components.append(1.0)
|
|
return (float(components[0]), float(components[1]), float(components[2]), float(components[3]))
|
|
|
|
|
|
def _resolve_configured_artifact_path(
|
|
raw_value: object,
|
|
*,
|
|
base_path: Path,
|
|
default_path: Path | None,
|
|
expected_suffixes: tuple[str, ...],
|
|
field_name: str,
|
|
) -> Path:
|
|
if raw_value is None:
|
|
if default_path is None:
|
|
raise CadSourceError(f"{_relative_to_repo(base_path)} {field_name} is required")
|
|
resolved = default_path.resolve()
|
|
else:
|
|
if not isinstance(raw_value, str) or not raw_value.strip():
|
|
raise CadSourceError(f"{_relative_to_repo(base_path)} {field_name} must be a non-empty string")
|
|
value = raw_value.strip()
|
|
if "\\" in value:
|
|
raise CadSourceError(f"{_relative_to_repo(base_path)} {field_name} must use POSIX '/' separators")
|
|
pure = PurePosixPath(value)
|
|
if pure.is_absolute() or any(part in {"", "."} for part in pure.parts):
|
|
raise CadSourceError(f"{_relative_to_repo(base_path)} {field_name} must be relative")
|
|
resolved = (base_path.parent.resolve() / Path(*pure.parts)).resolve()
|
|
suffix = resolved.suffix.lower()
|
|
if suffix not in expected_suffixes:
|
|
joined = " or ".join(expected_suffixes)
|
|
raise CadSourceError(f"{_relative_to_repo(base_path)} {field_name} must end in {joined}")
|
|
return resolved
|
|
|
|
|
|
def _required_output(raw_value: str | None, *, script_path: Path, field_name: str) -> str:
|
|
if raw_value is None:
|
|
raise CadSourceError(f"{_relative_to_repo(script_path)} {field_name} is required")
|
|
return raw_value
|
|
|
|
|
|
def _source_label(source: CadSource) -> str:
|
|
if source.script_path is not None:
|
|
return _relative_to_repo(source.script_path)
|
|
return _relative_to_repo(source.source_path)
|
|
|
|
|
|
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()
|