chore: snapshot workflow migration progress

render-worker/scripts/export_step_to_usd.py

@@ -29,11 +29,14 @@ import argparse
 import hashlib
 import json
 import math
+import os
 import re
 import sys
 import traceback
 from pathlib import Path
 
+from _blender_materials import build_mat_map_lower, lookup_material_name
+
 
 # ── CLI ───────────────────────────────────────────────────────────────────────
 
@@ -574,9 +577,8 @@ def _author_xcaf_to_usd(
         mat_usd_path = f"/Root/Looks/{mat_prim_name}"
         if not stage.GetPrimAtPath(mat_usd_path):
             UsdShade.Material.Define(stage, mat_usd_path)
-        UsdShade.MaterialBindingAPI(mesh.GetPrim()).Bind(
-            UsdShade.Material(stage.GetPrimAtPath(mat_usd_path))
-        )
+        binding_api = UsdShade.MaterialBindingAPI.Apply(mesh.GetPrim())
+        binding_api.Bind(UsdShade.Material(stage.GetPrimAtPath(mat_usd_path)))
 
         manifest_parts.append({
             "part_key": part_key,
@@ -968,57 +970,15 @@ def _prim_name(name: str) -> str:
     return safe or "unnamed"
 
 
-# ── Material map lookup (mirrors _blender_materials.build_mat_map_lower) ─────
-
-def _build_mat_map_lower(material_map: dict) -> dict:
-    """Build a lowercased material_map with AF-stripped and slug variants.
-
-    Same normalization as _blender_materials.build_mat_map_lower() so that
-    source_name → canonical material name lookup works consistently.
-    """
-    mat_map_lower: dict = {}
-    for k, v in material_map.items():
-        kl = k.lower().strip()
-        mat_map_lower[kl] = v
-        # Slug variant: replace non-alphanumeric with '_' (same as _generate_part_key)
-        slug_key = re.sub(r'[^a-z0-9]+', '_', kl).strip('_')
-        if slug_key and slug_key != kl:
-            mat_map_lower.setdefault(slug_key, v)
-        # Strip OCC assembly-frame suffixes: _AF0, _AF0_1, _AF0_1_AF0, etc.
-        stripped = re.sub(r'(_af\d+(_\d+)?)+$', '', kl)
-        if stripped != kl:
-            mat_map_lower.setdefault(stripped, v)
-            slug_stripped = re.sub(r'[^a-z0-9]+', '_', stripped).strip('_')
-            if slug_stripped and slug_stripped != stripped:
-                mat_map_lower.setdefault(slug_stripped, v)
-    return mat_map_lower
-
-
 def _lookup_material(source_name: str, part_key: str, mat_map_lower: dict) -> str | None:
     """Look up canonical material name for a part, trying multiple key variants."""
-    if not mat_map_lower:
-        return None
-    # Try source_name (lowered)
-    sn = source_name.lower().strip()
-    if sn in mat_map_lower:
-        return mat_map_lower[sn]
-    # Try AF-stripped source_name
-    stripped = re.sub(r'(_af\d+(_\d+)?)+$', '', sn, flags=re.IGNORECASE)
-    if stripped != sn and stripped in mat_map_lower:
-        return mat_map_lower[stripped]
-    # Try slug of source_name (matches part_key generation logic)
-    slug = re.sub(r'[^a-z0-9]+', '_', sn).strip('_')
-    if slug and slug in mat_map_lower:
-        return mat_map_lower[slug]
-    # Try part_key directly
-    pk = part_key.lower().strip()
-    if pk in mat_map_lower:
-        return mat_map_lower[pk]
-    # Prefix fallback: longest key that starts with or is started by part_key
-    for key in sorted(mat_map_lower.keys(), key=len, reverse=True):
-        if len(key) >= 5 and len(pk) >= 5 and (pk.startswith(key) or key.startswith(pk)):
-            return mat_map_lower[key]
-    return None
+    return lookup_material_name(source_name, mat_map_lower, part_key)
+
+
+def _atomic_output_path(output_path: Path) -> Path:
+    return output_path.with_name(
+        f".{output_path.stem}.{os.getpid()}.tmp{output_path.suffix}"
+    )
 
 
 # ── Main ──────────────────────────────────────────────────────────────────────
@@ -1027,7 +987,7 @@ def main() -> None:
     args = parse_args()
     color_map: dict = json.loads(args.color_map)
     raw_material_map: dict = json.loads(args.material_map)
-    mat_map_lower = _build_mat_map_lower(raw_material_map) if raw_material_map else {}
+    mat_map_lower = build_mat_map_lower(raw_material_map) if raw_material_map else {}
     if mat_map_lower:
         print(f"Material map: {len(raw_material_map)} entries ({len(mat_map_lower)} with variants)")
 
@@ -1165,7 +1125,14 @@ def main() -> None:
             print(f"WARNING: palette colors failed (non-fatal): {exc}", file=sys.stderr)
 
     # ── Create USD stage ──────────────────────────────────────────────────
-    stage = Usd.Stage.CreateNew(str(output_path))
+    temp_output_path = _atomic_output_path(output_path)
+    try:
+        if temp_output_path.exists():
+            temp_output_path.unlink()
+    except OSError:
+        pass
+
+    stage = Usd.Stage.CreateNew(str(temp_output_path))
     UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)
     UsdGeom.SetStageMetersPerUnit(stage, 0.001)  # mm; Blender handles m conversion on import
 
@@ -1206,7 +1173,19 @@ def main() -> None:
 
     n_parts = counters["n_parts"]
     n_empty = counters["n_empty"]
-    stage.Save()
+    try:
+        stage.Save()
+        if temp_output_path.exists():
+            os.chmod(temp_output_path, 0o664)
+        os.replace(temp_output_path, output_path)
+        os.chmod(output_path, 0o664)
+    except Exception:
+        try:
+            if temp_output_path.exists():
+                temp_output_path.unlink()
+        except OSError:
+            pass
+        raise
 
     sz = output_path.stat().st_size // 1024 if output_path.exists() else 0
     n_mat_assigned = sum(1 for p in manifest_parts if p.get("canonical_material"))