Recursively export xcresult screenshot attachments

Author: shai-almog

scripts/run-ios-ui-tests.sh

@@ -174,7 +174,186 @@ rm -rf "$EXPORT_DIR"
 mkdir -p "$EXPORT_DIR"
 
 ri_log "Exporting screenshot attachments from $RESULT_BUNDLE"
-if ! xcrun xcresulttool export --legacy --path "$RESULT_BUNDLE" --type file --output-path "$EXPORT_DIR" >/dev/null; then
+if ! python3 - "$RESULT_BUNDLE" "$EXPORT_DIR" <<'PY'
+import json
+import os
+import re
+import subprocess
+import sys
+from collections import deque
+
+
+def ri_log(message: str) -> None:
+    print(f"[run-ios-ui-tests] {message}", file=sys.stderr)
+
+
+def run_xcresult(args):
+    cmd = ["xcrun", "xcresulttool", *args]
+    try:
+        result = subprocess.run(
+            cmd,
+            check=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+    except subprocess.CalledProcessError as exc:
+        ri_log(
+            "xcresulttool command failed: {}\n{}".format(
+                " ".join(cmd), exc.stderr.decode("utf-8", "ignore").strip()
+            )
+        )
+        raise
+    return result.stdout.decode("utf-8")
+
+
+def get_json(bundle_path, object_id=None):
+    args = ["get", "--path", bundle_path, "--format", "json"]
+    if object_id:
+        args.extend(["--id", object_id])
+    output = run_xcresult(args)
+    return json.loads(output or "{}")
+
+
+def collect_nodes(node):
+    attachments = []
+    refs = []
+    stack = [node]
+    while stack:
+        current = stack.pop()
+        if isinstance(current, dict):
+            attachment_block = current.get("attachments")
+            if isinstance(attachment_block, dict):
+                for item in attachment_block.get("_values", []):
+                    if isinstance(item, dict):
+                        attachments.append(item)
+            for key, value in current.items():
+                if key.endswith("Ref") and isinstance(value, dict):
+                    ref_id = value.get("id")
+                    if isinstance(ref_id, str) and ref_id:
+                        refs.append(ref_id)
+                elif key.endswith("Refs") and isinstance(value, dict):
+                    for entry in value.get("_values", []):
+                        if isinstance(entry, dict):
+                            ref_id = entry.get("id")
+                            if isinstance(ref_id, str) and ref_id:
+                                refs.append(ref_id)
+                if isinstance(value, (dict, list)):
+                    stack.append(value)
+        elif isinstance(current, list):
+            stack.extend(current)
+    return attachments, refs
+
+
+def is_image_attachment(attachment):
+    uti = (attachment.get("uniformTypeIdentifier") or "").lower()
+    filename = (attachment.get("filename") or attachment.get("name") or "").lower()
+    if filename.endswith((".png", ".jpg", ".jpeg")):
+        return True
+    if "png" in uti or "jpeg" in uti:
+        return True
+    return False
+
+
+def sanitize_filename(name):
+    safe = re.sub(r"[^A-Za-z0-9_.-]", "_", name)
+    return safe or "attachment"
+
+
+def ensure_extension(name, uti):
+    base, ext = os.path.splitext(name)
+    if ext:
+        return name
+    uti = uti.lower()
+    if "png" in uti:
+        return f"{name}.png"
+    if "jpeg" in uti:
+        return f"{name}.jpg"
+    return name
+
+
+def export_attachment(bundle_path, attachment, destination_dir, used_names):
+    payload = attachment.get("payloadRef") or {}
+    attachment_id = payload.get("id")
+    if not attachment_id:
+        return
+    name = attachment.get("filename") or attachment.get("name") or attachment_id
+    name = ensure_extension(name, attachment.get("uniformTypeIdentifier") or "")
+    name = sanitize_filename(name)
+    candidate = name
+    counter = 1
+    while candidate in used_names:
+        base, ext = os.path.splitext(name)
+        candidate = f"{base}_{counter}{ext}"
+        counter += 1
+    used_names.add(candidate)
+    output_path = os.path.join(destination_dir, candidate)
+    run_xcresult(
+        [
+            "export",
+            "--legacy",
+            "--path",
+            bundle_path,
+            "--id",
+            attachment_id,
+            "--type",
+            "file",
+            "--output-path",
+            output_path,
+        ]
+    )
+    ri_log(f"Exported attachment {candidate}")
+
+
+def main():
+    if len(sys.argv) != 3:
+        ri_log("Expected bundle path and destination directory arguments")
+        return 1
+
+    bundle_path, destination_dir = sys.argv[1:3]
+    os.makedirs(destination_dir, exist_ok=True)
+
+    root = get_json(bundle_path)
+    attachments, refs = collect_nodes(root)
+    queue = deque(refs)
+    seen_refs = set()
+    seen_attachment_ids = set()
+
+    def handle_attachments(items):
+        for attachment in items:
+            payload = attachment.get("payloadRef") or {}
+            attachment_id = payload.get("id")
+            if not attachment_id or attachment_id in seen_attachment_ids:
+                continue
+            if not is_image_attachment(attachment):
+                continue
+            seen_attachment_ids.add(attachment_id)
+            export_attachment(bundle_path, attachment, destination_dir, exported_names)
+
+    exported_names = set()
+    handle_attachments(attachments)
+
+    while queue:
+        ref_id = queue.popleft()
+        if ref_id in seen_refs:
+            continue
+        seen_refs.add(ref_id)
+        data = get_json(bundle_path, ref_id)
+        items, nested_refs = collect_nodes(data)
+        handle_attachments(items)
+        for nested in nested_refs:
+            if nested not in seen_refs:
+                queue.append(nested)
+
+    if not exported_names:
+        ri_log("No screenshot attachments were exported from xcresult bundle")
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
+PY
+then
   ri_log "xcresulttool export failed"
   exit 11
 fi