VidSimplify / manimator /utils /code_postprocessor.py
Adityahulk
Restoring repo state for deployment
6fc3143
"""
Code Post-Processor for Manim
Fixes common issues in AI-generated Manim code, particularly around
SurroundingRectangle with indexed MathTex elements.
"""
import re
from typing import List
def fix_surrounding_rectangles(code: str) -> str:
"""
Fix or remove SurroundingRectangle calls that use indexed access to MathTex.
Problem: AI often generates code like:
SurroundingRectangle(equation[0][5], ...)
This doesn't work reliably because MathTex indexing is unpredictable.
Solution: Comment out these problematic lines with an explanation.
Args:
code: Raw generated Manim code
Returns:
Fixed code with problematic SurroundingRectangle calls commented out
"""
lines = code.split('\n')
fixed_lines = []
skip_next = False # Initialize before the loop
for line in lines:
# Pattern: SurroundingRectangle with indexed access like equation[0][5]
# Matches: SurroundingRectangle(something[X][Y], ...) or SurroundingRectangle(something[X:Y], ...)
if 'SurroundingRectangle(' in line:
# Check if it has indexed or sliced access
# Pattern: variable[number] or variable[number:number] inside SurroundingRectangle()
if re.search(r'SurroundingRectangle\([^,\)]*\[\d+\]', line) or \
re.search(r'SurroundingRectangle\([^,\)]*\[\d+:\d*\]', line):
# Replace with empty VGroup to prevent NameError in subsequent animations
# while avoiding the rendering crash
indent = len(line) - len(line.lstrip())
# Extract variable name if present (e.g. "box = SurroundingRectangle(...)")
var_name = ""
if "=" in line:
var_name = line.split("=")[0].strip()
replacement = ' ' * indent + f'{var_name} = VGroup() # Replaced invalid SurroundingRectangle'
else:
# If no assignment, just comment it out as it won't be referenced
replacement = ' ' * indent + '# ' + line.lstrip() + ' # Auto-disabled: indexed SurroundingRectangle'
fixed_lines.append(replacement)
else:
fixed_lines.append(line)
else:
fixed_lines.append(line)
return '\n'.join(fixed_lines)
def remove_problematic_indexing(code: str) -> str:
"""
More aggressive approach: Remove entire blocks that use indexed MathTex highlighting.
This removes the entire voiceover block if it contains indexed SurroundingRectangle.
"""
# For now, use the commenting approach which is safer
return fix_surrounding_rectangles(code)
def fix_undefined_colors(code: str) -> str:
"""
Fix undefined color constants by replacing them with valid Manim colors.
Common issues:
- ORANGE_A, ORANGE_B, etc. -> ORANGE
- RED_A, RED_B, etc. -> RED
- Similar patterns for other colors
Args:
code: Raw generated code
Returns:
Code with undefined colors replaced
"""
# Color mappings: undefined variants -> standard colors
color_replacements = {
# Orange variants
r'\bORANGE_[A-Z]\b': 'ORANGE',
# Red variants
r'\bRED_[A-Z]\b': 'RED',
# Blue variants
r'\bBLUE_[A-Z]\b': 'BLUE',
# Green variants
r'\bGREEN_[A-Z]\b': 'GREEN',
# Yellow variants
r'\bYELLOW_[A-Z]\b': 'YELLOW',
# Purple variants
r'\bPURPLE_[A-Z]\b': 'PURPLE',
# Pink variants
r'\bPINK_[A-Z]\b': 'PINK',
# Teal variants
r'\bTEAL_[A-Z]\b': 'TEAL',
# Gray variants
r'\bGRAY_[A-Z]\b': 'GRAY',
}
for pattern, replacement in color_replacements.items():
code = re.sub(pattern, replacement, code)
return code
def clean_duplicate_imports(code: str) -> str:
"""
Remove duplicate and incorrect imports.
Keep only the correct imports for our architecture.
"""
lines = code.split('\n')
cleaned_lines = []
seen_imports = set()
# Imports to remove (old/wrong patterns)
bad_imports = [
'from manim_voiceover import VoiceoverScene',
'from manimator.services import ElevenLabsService',
'from manim_voiceover.services.gtts import GTTSService',
'from manim_voiceover.services.elevenlabs import ElevenLabsService',
]
for line in lines:
stripped = line.strip()
# Skip bad imports
if any(bad_import in line for bad_import in bad_imports):
continue
# Track and deduplicate good imports
if stripped.startswith('from ') or stripped.startswith('import '):
if stripped not in seen_imports:
seen_imports.add(stripped)
cleaned_lines.append(line)
else:
cleaned_lines.append(line)
return '\n'.join(cleaned_lines)
def fix_invalid_manim_parameters(code: str) -> str:
"""
Fix invalid Manim parameters that cause runtime errors.
Common issues:
- rounded_corners parameter in Rectangle (doesn't exist)
- corner_radius parameter in Rectangle (doesn't exist in older Manim versions)
- Invalid parameter names
"""
# Remove corner_radius from Rectangle/RoundedRectangle calls (line by line)
lines = code.split('\n')
fixed_lines = []
for line in lines:
# Skip lines that are just the parameter
if 'corner_radius' in line and '=' in line:
# Check if this is a standalone parameter line
stripped = line.strip()
if stripped.startswith('corner_radius=') or stripped.startswith('rounded_corners='):
# Skip this line entirely
continue
# Remove inline corner_radius/rounded_corners parameters
line = re.sub(r',\s*corner_radius\s*=\s*[^,)]+', '', line)
line = re.sub(r',\s*rounded_corners\s*=\s*[^,)]+', '', line)
# Remove invalid scale_tips parameter (often hallucinated for scale() method)
if 'scale_tips' in line and '=' in line:
stripped = line.strip()
if stripped.startswith('scale_tips='):
continue
# Handle inline scale_tips removal (both with and without leading comma)
line = re.sub(r',\s*scale_tips\s*=\s*[^,)]+', '', line)
line = re.sub(r'\bscale_tips\s*=\s*[^,)]+\s*,?', '', line)
fixed_lines.append(line)
return '\n'.join(fixed_lines)
def fix_visual_leaks(code: str) -> str:
"""
Prevent 'visual memory leaks' where reassigning a variable leaves the old Mobject on screen.
Also enforces 'Title Exclusivity' to prevent multiple titles from stacking.
Problem 1 (Variable Reuse):
text = Text("A")
self.play(Write(text))
text = Text("B") # "A" is still on screen!
Problem 2 (Title Stacking):
title1 = Text("Intro")
self.play(Write(title1))
title2 = Text("Chapter 1") # "Intro" is still on screen!
Solution:
Inject `self.remove(var)` checks.
"""
lines = code.split('\n')
fixed_lines = []
# Regex to detect Mobject assignments
# Matches: var = ClassName(...)
# We broaden this to catch any Capitalized class instantiation to be safer
assignment_pattern = re.compile(r'^\s*([a-zA-Z_]\w*)\s*=\s*([A-Z]\w*)\(')
# Track variables that have been assigned
assigned_vars = set()
# Track variables that look like titles
active_titles = []
for line in lines:
match = assignment_pattern.match(line)
if match:
var_name = match.group(1)
class_name = match.group(2)
indent = len(line) - len(line.lstrip())
indent_str = ' ' * indent
# 1. Handle Variable Reuse (Same variable name)
# We inject this for EVERY assignment to handle loops correctly.
# In a loop, the line `t = Text(...)` appears once but runs multiple times.
# By injecting the check, we ensure the previous iteration's object is removed.
removal = f"{indent_str}if '{var_name}' in locals(): self.remove({var_name})"
fixed_lines.append(removal)
# 2. Handle Title Exclusivity (Different variable names, but both are titles)
# Check if this looks like a title (contains 'title', 'header', 'heading')
# But exclude 'subtitle' or 'sub_title'
is_title = re.search(r'title|header|heading', var_name, re.IGNORECASE) and \
not re.search(r'sub', var_name, re.IGNORECASE)
if is_title:
# If we are creating a new title, remove ALL previous titles
for old_title in active_titles:
# Don't remove if it's the same variable (handled above)
if old_title != var_name:
removal = f"{indent_str}if '{old_title}' in locals(): self.remove({old_title})"
fixed_lines.append(removal)
# Add this to active titles (if not already there)
if var_name not in active_titles:
active_titles.append(var_name)
assigned_vars.add(var_name)
fixed_lines.append(line)
else:
fixed_lines.append(line)
return '\n'.join(fixed_lines)
def post_process_code(code: str) -> str:
"""
Main entry point for code post-processing.
Applies all fixes to AI-generated Manim code.
Args:
code: Raw generated code
Returns:
Cleaned and fixed code
"""
# Check if we need to add header (before making changes)
has_undefined_colors = bool(re.search(r'\b(ORANGE|RED|BLUE|GREEN|YELLOW|PURPLE|PINK|TEAL|GRAY)_[A-Z]\b', code))
# Apply fixes in order
code = clean_duplicate_imports(code)
code = fix_undefined_colors(code)
code = fix_invalid_manim_parameters(code)
code = fix_invalid_manim_parameters(code)
code = fix_surrounding_rectangles(code)
code = fix_visual_leaks(code)
# Add header comment explaining post-processing
header = """# NOTE: This code has been automatically post-processed to fix common issues.
# Indexed SurroundingRectangle calls have been disabled as they don't reliably
# highlight the intended equation parts in MathTex objects.
# Undefined color constants have been replaced with standard Manim colors.
# Invalid Manim parameters have been removed or corrected.
"""
# Only add header if we actually made changes
if '# Auto-disabled:' in code or has_undefined_colors or 'rounded_corners' in code:
code = header + code
return code
def validate_code_structure(code: str) -> List[str]:
"""
Validate the generated code for common issues.
Returns:
List of warning messages (empty if no issues)
"""
warnings = []
# Check for common issues
if 'SurroundingRectangle(' in code:
if re.search(r'SurroundingRectangle\([^,\)]*\[\d+\]', code):
warnings.append("Code contains indexed SurroundingRectangle calls (will be auto-fixed)")
if 'from manim_voiceover.services.gtts import GTTSService' in code:
warnings.append("Code still uses deprecated GTTSService (should use ElevenLabsService)")
return warnings