executor.py

"""Render generated Manim scenes in an isolated subprocess."""

from __future__ import annotations

import shutil
import subprocess
import tempfile
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
import re
import sys
from typing import Tuple


OUTPUT_DIR = Path("./output")
SCENE_CLASS_PATTERN = re.compile(r"class\s+(\w+)\s*\(\s*\w*Scene\w*\s*\):")


@dataclass
class RenderResult:
    success: bool
    video_path: str | None
    info: str
    errors: str


def find_scene_name(code: str) -> str:
    """Extract the first Manim Scene class name from generated Python code."""
    match = SCENE_CLASS_PATTERN.search(code)
    return match.group(1) if match else "MainScene"


def _manim_command_prefix() -> list[str]:
    """Return a command prefix for Manim.

    Hugging Face Spaces occasionally installs the Python package correctly
    without exposing the `manim` console script on PATH. `python -m manim`
    still works in that case, so keep it as the final fallback.
    """
    for command in ("manim", "manimce"):
        if shutil.which(command):
            return [command]
    return [sys.executable, "-m", "manim"]


def render_code(
    code: str,
    timeout_seconds: int = 600,
    output_dir: Path = OUTPUT_DIR,
    quality: str = "-ql",
) -> RenderResult:
    """Render Manim code and return a structured result.

    This mirrors the robust approach used by manim-trainer: use a temporary
    script and a temporary media directory to avoid collisions, then write only
    the final MP4 into the app's output folder.
    """
    if not code or not code.strip():
        return RenderResult(False, None, "", "No Manim code was provided.")

    output_dir.mkdir(parents=True, exist_ok=True)
    scene_name = find_scene_name(code)
    video_name = f"{scene_name}_{uuid.uuid4().hex[:12]}_{int(time.time())}"
    final_video_base = output_dir / video_name

    temp_script_name = ""
    try:
        with tempfile.NamedTemporaryFile(
            mode="w",
            suffix=".py",
            prefix="temp_scene_",
            encoding="utf-8",
            delete=False,
        ) as temp_script:
            temp_script.write(code)
            temp_script.flush()
            temp_script_name = temp_script.name

        with tempfile.TemporaryDirectory(prefix="scivisual_manim_") as temp_media_dir:
            command = _manim_command_prefix() + [
                temp_script_name,
                scene_name,
                quality,
                "--format=mp4",
                "--media_dir",
                temp_media_dir,
                "-o",
                str(final_video_base.absolute()),
            ]
            result = subprocess.run(
                command,
                capture_output=True,
                text=True,
                timeout=timeout_seconds,
                check=False,
            )

        stdout = result.stdout.strip()
        stderr = result.stderr.strip()
        if result.returncode != 0:
            error_log = stderr or stdout or f"Manim failed with return code {result.returncode}."
            return RenderResult(False, None, stdout, error_log)

        final_video_path = final_video_base.with_suffix(".mp4")
        if final_video_path.exists():
            return RenderResult(True, str(final_video_path), stdout, stderr)

        return RenderResult(
            False,
            None,
            stdout,
            f"Manim reported success, but no MP4 was found at {final_video_path}. STDERR: {stderr}",
        )
    except FileNotFoundError as exc:
        return RenderResult(
            False,
            None,
            "",
            "Manim could not be launched. Ensure the `manim` Python package is installed. "
            f"Attempted fallback command: `{sys.executable} -m manim`. Details: {exc}",
        )
    except subprocess.TimeoutExpired as exc:
        stdout = exc.stdout or ""
        stderr = exc.stderr or ""
        return RenderResult(
            False,
            None,
            stdout,
            f"Rendering timed out after {timeout_seconds} seconds.\n\nSTDERR:\n{stderr}",
        )
    except Exception as exc:  # pragma: no cover - runtime safety guard
        return RenderResult(False, None, "", f"Unexpected render failure: {type(exc).__name__}: {exc}")
    finally:
        if temp_script_name:
            Path(temp_script_name).unlink(missing_ok=True)


def render_manim_scene(code: str, timeout_seconds: int = 600) -> Tuple[bool, str]:
    """Compatibility wrapper returning the tuple shape used by the app."""
    result = render_code(code, timeout_seconds=timeout_seconds)
    if result.success and result.video_path:
        return True, result.video_path
    return False, result.errors or result.info