mirror of
https://github.com/3b1b/manim.git
synced 2025-04-13 09:47:07 +00:00
152 lines
4.7 KiB
Python
152 lines
4.7 KiB
Python
import inspect
|
|
import pyperclip
|
|
import os
|
|
|
|
from IPython.core.getipython import get_ipython
|
|
from IPython.terminal import pt_inputhooks
|
|
from IPython.terminal.embed import InteractiveShellEmbed
|
|
|
|
from manimlib.animation.fading import VFadeInThenOut
|
|
from manimlib.constants import RED
|
|
from manimlib.mobject.mobject import Mobject
|
|
from manimlib.mobject.frame import FullScreenRectangle
|
|
from manimlib.module_loader import ModuleLoader
|
|
|
|
|
|
def interactive_scene_embed(scene):
|
|
scene.stop_skipping()
|
|
scene.update_frame(force_draw=True)
|
|
|
|
shell = get_ipython_shell_for_embedded_scene(scene)
|
|
enable_gui(shell, scene)
|
|
ensure_frame_update_post_cell(shell, scene)
|
|
ensure_flash_on_error(shell, scene)
|
|
|
|
# Launch shell
|
|
shell()
|
|
|
|
|
|
def get_ipython_shell_for_embedded_scene(scene):
|
|
"""
|
|
Create embedded IPython terminal configured to have access to
|
|
the local namespace of the caller
|
|
"""
|
|
# Triple back should take us to the context in a user's scene definition
|
|
# which is calling "self.embed"
|
|
caller_frame = inspect.currentframe().f_back.f_back.f_back
|
|
|
|
# Update the module's namespace to include local variables
|
|
module = ModuleLoader.get_module(caller_frame.f_globals["__file__"])
|
|
module.__dict__.update(caller_frame.f_locals)
|
|
module.__dict__.update(get_shortcuts(scene))
|
|
|
|
return InteractiveShellEmbed(
|
|
user_module=module,
|
|
display_banner=False,
|
|
xmode=scene.embed_exception_mode
|
|
)
|
|
|
|
|
|
def get_shortcuts(scene):
|
|
"""
|
|
A few custom shortcuts useful to have in the interactive shell namespace
|
|
"""
|
|
return dict(
|
|
play=scene.play,
|
|
wait=scene.wait,
|
|
add=scene.add,
|
|
remove=scene.remove,
|
|
clear=scene.clear,
|
|
focus=scene.focus,
|
|
save_state=scene.save_state,
|
|
reload=scene.reload,
|
|
undo=scene.undo,
|
|
redo=scene.redo,
|
|
i2g=scene.i2g,
|
|
i2m=scene.i2m,
|
|
checkpoint_paste=scene.checkpoint_paste,
|
|
)
|
|
|
|
|
|
def enable_gui(shell, scene):
|
|
"""Enables gui interactions during the embed"""
|
|
def inputhook(context):
|
|
while not context.input_is_ready():
|
|
if not scene.is_window_closing():
|
|
scene.update_frame(dt=0)
|
|
if scene.is_window_closing():
|
|
shell.ask_exit()
|
|
|
|
pt_inputhooks.register("manim", inputhook)
|
|
shell.enable_gui("manim")
|
|
|
|
|
|
def ensure_frame_update_post_cell(shell, scene):
|
|
"""Ensure the scene updates its frame after each ipython cell"""
|
|
def post_cell_func(*args, **kwargs):
|
|
if not scene.is_window_closing():
|
|
scene.update_frame(dt=0, force_draw=True)
|
|
|
|
shell.events.register("post_run_cell", post_cell_func)
|
|
|
|
|
|
def ensure_flash_on_error(shell, scene):
|
|
"""Flash border, and potentially play sound, on exceptions"""
|
|
def custom_exc(shell, etype, evalue, tb, tb_offset=None):
|
|
# Show the error don't just swallow it
|
|
shell.showtraceback((etype, evalue, tb), tb_offset=tb_offset)
|
|
if scene.embed_error_sound:
|
|
os.system("printf '\a'")
|
|
rect = FullScreenRectangle().set_stroke(RED, 30).set_fill(opacity=0)
|
|
rect.fix_in_frame()
|
|
scene.play(VFadeInThenOut(rect, run_time=0.5))
|
|
|
|
shell.set_custom_exc((Exception,), custom_exc)
|
|
|
|
|
|
class CheckpointManager:
|
|
checkpoint_states: dict[str, list[tuple[Mobject, Mobject]]] = dict()
|
|
|
|
def checkpoint_paste(self, scene):
|
|
"""
|
|
Used during interactive development to run (or re-run)
|
|
a block of scene code.
|
|
|
|
If the copied selection starts with a comment, this will
|
|
revert to the state of the scene the first time this function
|
|
was called on a block of code starting with that comment.
|
|
"""
|
|
shell = get_ipython()
|
|
if shell is None:
|
|
return
|
|
|
|
code_string = pyperclip.paste()
|
|
|
|
checkpoint_key = self.get_leading_comment(code_string)
|
|
self.handle_checkpoint_key(scene, checkpoint_key)
|
|
shell.run_cell(code_string)
|
|
|
|
@staticmethod
|
|
def get_leading_comment(code_string: str):
|
|
leading_line = code_string.partition("\n")[0].lstrip()
|
|
if leading_line.startswith("#"):
|
|
return leading_line
|
|
return None
|
|
|
|
def handle_checkpoint_key(self, scene, key: str):
|
|
if key is None:
|
|
return
|
|
elif key in self.checkpoint_states:
|
|
# Revert to checkpoint
|
|
scene.restore_state(self.checkpoint_states[key])
|
|
|
|
# Clear out any saved states that show up later
|
|
all_keys = list(self.checkpoint_states.keys())
|
|
index = all_keys.index(key)
|
|
for later_key in all_keys[index + 1:]:
|
|
self.checkpoint_states.pop(later_key)
|
|
else:
|
|
self.checkpoint_states[key] = scene.get_state()
|
|
|
|
def clear_checkpoints(self):
|
|
self.checkpoint_states = dict()
|