manim_executor.py

"""
Manim code execution and rendering handler
"""

import os
import tempfile
import subprocess
import sys
import shutil
from pathlib import Path
import re
import traceback

class ManimExecutor:
    def __init__(self):
        self.temp_dir = tempfile.mkdtemp(prefix="manim_")
        self.output_dir = os.path.join(self.temp_dir, "media")
    
    def render_animation(self, code, quality="medium_quality", output_format="mp4"):
        """
        Execute Manim code and render animation
        
        Args:
            code (str): Python code containing Manim scene
            quality (str): Rendering quality (low_quality, medium_quality, high_quality)
            output_format (str): Output format (mp4, gif)
        
        Returns:
            dict: Result containing success status, output_file path, or error message
        """
        try:
            # Validate code
            validation_result = self._validate_code(code)
            if not validation_result["valid"]:
                return {"success": False, "error": validation_result["error"]}
            
            # Extract scene class name
            scene_name = self._extract_scene_name(code)
            if not scene_name:
                return {"success": False, "error": "No Scene class found in the code. Please define a class that inherits from Scene."}
            
            # Write code to temporary file
            code_file = os.path.join(self.temp_dir, "animation.py")
            with open(code_file, "w", encoding="utf-8") as f:
                f.write(code)
            
            # Prepare manim command
            # Map quality settings to Manim's single-letter quality flags
            quality_map = {
                "low_quality": "l",
                "medium_quality": "m",
                "high_quality": "h"
            }
            quality_letter = quality_map.get(quality, "m")
            
            cmd = [
                sys.executable, "-m", "manim", "render",
                "-q", quality_letter,
                "--format", output_format,
                "-o", f"output.{output_format}",
                code_file,
                scene_name
            ]
            
            # Execute manim command
            result = subprocess.run(
                cmd,
                cwd=self.temp_dir,
                capture_output=True,
                text=True,
                timeout=60  # 60 second timeout
            )
            
            if result.returncode != 0:
                error_msg = self._parse_manim_error(result.stderr)
                return {"success": False, "error": error_msg}
            
            # Find output file
            output_file = self._find_output_file(output_format)
            if not output_file:
                return {"success": False, "error": "Animation was rendered but output file not found."}
            
            return {"success": True, "output_file": output_file}
        
        except subprocess.TimeoutExpired:
            return {"success": False, "error": "Rendering timed out. Please try a simpler animation or reduce the complexity."}
        except Exception as e:
            return {"success": False, "error": f"Execution error: {str(e)}"}
    
    def _validate_code(self, code):
        """Validate the Python code for basic syntax and security"""
        try:
            # Check for potentially dangerous imports/functions
            dangerous_patterns = [
                r'import\s+os',
                r'from\s+os\s+import',
                r'import\s+subprocess',
                r'import\s+sys',
                r'open\s*\(',
                r'exec\s*\(',
                r'eval\s*\(',
                r'__import__',
                r'file\s*\(',
                r'input\s*\(',
            ]
            
            for pattern in dangerous_patterns:
                if re.search(pattern, code, re.IGNORECASE):
                    return {
                        "valid": False,
                        "error": f"Code contains potentially unsafe operations. Please remove: {pattern}"
                    }
            
            # Try to compile the code
            compile(code, '<string>', 'exec')
            
            return {"valid": True}
        
        except SyntaxError as e:
            return {
                "valid": False,
                "error": f"Syntax Error on line {e.lineno}: {e.msg}"
            }
        except Exception as e:
            return {
                "valid": False,
                "error": f"Code validation error: {str(e)}"
            }
    
    def _extract_scene_name(self, code):
        """Extract the Scene class name from the code"""
        # Look for class definitions that inherit from Scene
        pattern = r'class\s+(\w+)\s*\(\s*Scene\s*\):'
        matches = re.findall(pattern, code)
        
        if matches:
            return matches[0]  # Return the first Scene class found
        return None
    
    def _parse_manim_error(self, error_output):
        """Parse and simplify Manim error messages"""
        if not error_output:
            return "Unknown rendering error occurred."
        
        # Common error patterns and their simplified messages
        error_patterns = {
            r"ModuleNotFoundError": "Missing required module. Make sure all imports are correct.",
            r"NameError": "Variable or function name not found. Check your variable names.",
            r"AttributeError": "Object attribute error. Check your object methods and properties.",
            r"TypeError": "Type error. Check your function arguments and data types.",
            r"ValueError": "Invalid value provided. Check your numerical values and ranges.",
            r"IndentationError": "Indentation error. Check your code formatting.",
            r"SyntaxError": "Syntax error in your Python code.",
        }
        
        for pattern, message in error_patterns.items():
            if re.search(pattern, error_output, re.IGNORECASE):
                # Try to extract the specific error line
                lines = error_output.split('\n')
                for line in lines:
                    if 'line' in line.lower() and any(p in line for p in ['error', 'exception']):
                        return f"{message}\n\nDetailed error: {line.strip()}"
                return message
        
        # If no pattern matches, return a cleaned version of the error
        lines = error_output.split('\n')
        relevant_lines = []
        for line in lines:
            if any(keyword in line.lower() for keyword in ['error', 'exception', 'traceback']):
                relevant_lines.append(line.strip())
        
        if relevant_lines:
            return "Rendering error:\n" + "\n".join(relevant_lines[:5])  # Limit to first 5 relevant lines
        
        return f"Rendering failed with error:\n{error_output[:500]}"  # Limit to first 500 chars
    
    def _find_output_file(self, output_format):
        """Find the rendered output file"""
        # Manim typically saves files in media/videos or media/images
        search_dirs = [
            os.path.join(self.temp_dir, "media", "videos"),
            os.path.join(self.temp_dir, "media", "images"),
            os.path.join(self.temp_dir, "media"),
            self.temp_dir
        ]
        
        for search_dir in search_dirs:
            if os.path.exists(search_dir):
                for root, dirs, files in os.walk(search_dir):
                    for file in files:
                        if file.endswith(f'.{output_format}'):
                            return os.path.join(root, file)
        
        return None
    
    def cleanup(self):
        """Clean up temporary files"""
        try:
            if os.path.exists(self.temp_dir):
                shutil.rmtree(self.temp_dir)
        except Exception:
            pass  # Ignore cleanup errors
    
    def __del__(self):
        """Cleanup when object is destroyed"""
        self.cleanup()