scripts/ane_objectivec_probe.py

#!/usr/bin/env python3
"""Small local Objective-C/ANE runtime probe helper.

This keeps environment research local-only and non-destructive:
- compile-time checks for ObjC syntax and required frameworks;
- runtime checks for private ANE class presence (when classes are available);
- compact manifest output for local diagnostics.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional, Sequence


_ANE_CLASS_SET = (
    "_ANEClient",
    "_ANECompiler",
    "_ANERequest",
    "_ANEModel",
    "_ANEProgram",
    "_ANECompileOptions",
    "_ANEEngine",
    "_ANEInMemoryModelDescriptor",
    "_ANEData",
)


def _run_command(command: Sequence[str], *, timeout: int = 20, cwd: Optional[Path] = None) -> Dict[str, Any]:
    proc = subprocess.run(
        list(command),
        cwd=str(cwd) if cwd is not None else None,
        text=True,
        capture_output=True,
        timeout=timeout,
        check=False,
    )
    return {
        "command": " ".join(command),
        "code": int(proc.returncode),
        "stdout": (proc.stdout or "").strip(),
        "stderr": (proc.stderr or "").strip(),
    }


@dataclass(frozen=True)
class ProbeResult:
    manifest: Dict[str, Any]
    manifest_path: Optional[Path]


def _objc_source() -> str:
    class_checks = "\n".join(
        f'        payload[@"{name}"] = NSClassFromString(@"{name}") != Nil ? @YES : @NO;'
        for name in _ANE_CLASS_SET
    )
    template = """
#import <Foundation/Foundation.h>
#import <CoreML/CoreML.h>

int main(void) {{
    @autoreleasepool {{
        NSMutableDictionary *payload = [NSMutableDictionary new];
__CLASS_CHECKS__
        NSError *error = nil;
        NSData *json = [NSJSONSerialization dataWithJSONObject:payload options:0 error:&error];
        if (json == nil) {
            fprintf(stderr, "json-error:%s\\n", error ? [[error localizedDescription] UTF8String] : "unknown");
            return 2;
        }
        fwrite(json.bytes, 1, (unsigned long)json.length, stdout);
        return 0;
    }
}
"""
    return (
        template
        .replace("__CLASS_CHECKS__", class_checks)
        .replace("{{", "{")
        .replace("}}", "}")
    )


def _signature(payload: Dict[str, Any]) -> str:
    source = json.dumps(payload, sort_keys=True, default=str).encode("utf-8")
    return hashlib.sha256(source).hexdigest()[:16]


def run_objc_probe(out_json: Path) -> ProbeResult:
    out_json = Path(out_json)
    env = {
        "xcodebuild": Path("/usr/bin/xcodebuild").exists(),
        "xcrun": Path("/usr/bin/xcrun").exists(),
    }
    status: Dict[str, Any] = {
        "env": {
            "python": sys.executable,
            "cwd": str(Path.cwd()),
            "xcodebuild": os.environ.get("DEVELOPER_DIR", ""),
            "tools": env,
        },
        "checks": {},
    }

    sample_source = "#import <Foundation/Foundation.h>\n@interface Probe : NSObject @end\nint main(void){return 0;}\n"
    compile_syntax_cmd = [
        "xcrun",
        "clang",
        "-fsyntax-only",
        "-x",
        "objective-c",
        "-fobjc-arc",
        "-framework",
        "Foundation",
        "-framework",
        "CoreML",
        "-",
    ]
    compile_check = subprocess.run(
        compile_syntax_cmd,
        input=sample_source,
        text=True,
        capture_output=True,
        timeout=20,
        check=False,
    )
    status["checks"]["compile_syntax"] = {
        "command": " ".join(compile_syntax_cmd),
        "code": int(compile_check.returncode),
        "stdout": (compile_check.stdout or "").strip(),
        "stderr": (compile_check.stderr or "").strip(),
    }

    with tempfile.TemporaryDirectory() as workdir:
        tmp = Path(workdir)
        source = tmp / "ane_objc_probe.m"
        exe = tmp / "ane_objc_probe"
        source.write_text(_objc_source(), encoding="utf-8")

        compile_res = _run_command(
            [
                "xcrun",
                "clang",
                "-x",
                "objective-c",
                "-fobjc-arc",
                "-framework",
                "Foundation",
                "-framework",
                "CoreML",
                str(source),
                "-o",
                str(exe),
            ],
            timeout=30,
        )
        status["checks"]["runtime_compile"] = compile_res

        if compile_res["code"] == 0 and exe.exists():
            run_res = _run_command([str(exe)], timeout=20)
            status["checks"]["runtime_exec"] = run_res
            try:
                status["runtime_payload"] = json.loads(run_res.get("stdout", "{}"))
            except Exception:
                status["runtime_payload"] = {"error": "runtime-json-parse-failed", "raw": run_res.get("stdout", "")}
        else:
            status["runtime_payload"] = {
                "error": "runtime-compile-skipped",
            }

    status["checks"]["signature"] = _signature(status)
    out_json.parent.mkdir(parents=True, exist_ok=True)
    with out_json.open("w", encoding="utf-8") as fp:
        json.dump(status, fp, indent=2, default=str)
    return ProbeResult(manifest=status, manifest_path=out_json)


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = argparse.ArgumentParser(description="Local Objective-C/ANE probe.")
    parser.add_argument(
        "--out-json",
        default=str(Path.cwd() / "training" / "build_env" / "ane_objectivec_probe.json"),
        help="Where to write the JSON result.",
    )
    parser.add_argument(
        "--print-json",
        action="store_true",
        help="Print probe payload to stdout in JSON form.",
    )
    args = parser.parse_args(list(argv) if argv is not None else None)
    result = run_objc_probe(Path(args.out_json))
    if args.print_json:
        print(json.dumps(result.manifest, indent=2, default=str))
    else:
        print(f"ane_objectivec_probe_written={result.manifest_path}")
        print(f"signature={result.manifest.get('checks', {}).get('signature')}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())