Update InteractiveSceneEmbed (#2267)
Some checks are pending
docs / build up document and deploy (push) Waiting to run

* Comment tweak

* Directly print traceback

Since the shell.showtraceback is giving some issues

* Make InteracrtiveSceneEmbed into a class

This way it can keep track of it's internal shell; use of get_ipython has a finicky relationship with reloading.

* Move remaining checkpoint_paste logic into scene_embed.py

This involved making a few context managers for Scene: temp_record, temp_skip, temp_progress_bar, which seem useful in and of themselves.

* Change null key to be the empty string
This commit is contained in:
Grant Sanderson 2024-12-11 13:33:48 -06:00 committed by GitHub
parent eeb4fdf270
commit bafea89ac9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 158 additions and 149 deletions

View file

@ -146,7 +146,7 @@ class ModuleLoader:
if ignore_manimlib_modules and module.__name__.startswith("manimlib"):
return
if module.__name__.startswith("manimlib.config"):
# We don't want to reload global config
# We don't want to reload global manim_config
return
if not hasattr(module, "__dict__"):

View file

@ -5,6 +5,8 @@ import platform
import random
import time
from functools import wraps
from contextlib import contextmanager
from contextlib import ExitStack
from pyglet.window import key as PygletWindowKeys
@ -24,7 +26,7 @@ from manimlib.mobject.mobject import Mobject
from manimlib.mobject.mobject import Point
from manimlib.mobject.types.vectorized_mobject import VGroup
from manimlib.mobject.types.vectorized_mobject import VMobject
from manimlib.scene.scene_embed import interactive_scene_embed
from manimlib.scene.scene_embed import InteractiveSceneEmbed
from manimlib.scene.scene_embed import CheckpointManager
from manimlib.scene.scene_file_writer import SceneFileWriter
from manimlib.utils.dict_ops import merge_dicts_recursively
@ -122,8 +124,6 @@ class Scene(object):
self.time: float = 0
self.skip_time: float = 0
self.original_skipping_status: bool = self.skip_animations
self.checkpoint_states: dict[str, list[tuple[Mobject, Mobject]]] = dict()
self.checkpoint_manager: CheckpointManager = CheckpointManager()
self.undo_stack = []
self.redo_stack = []
@ -212,8 +212,10 @@ class Scene(object):
# Embed is only relevant for interactive development with a Window
return
self.show_animation_progress = show_animation_progress
self.stop_skipping()
self.update_frame(force_draw=True)
interactive_scene_embed(self)
InteractiveSceneEmbed(self).launch()
# End scene when exiting an embed
if close_scene_on_exit:
@ -678,41 +680,44 @@ class Scene(object):
self.undo_stack.append(self.get_state())
self.restore_state(self.redo_stack.pop())
def checkpoint_paste(
self,
skip: bool = False,
record: bool = False,
progress_bar: bool = True
):
"""
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.
"""
# Keep track of skipping and progress bar status
self.skip_animations = skip
@contextmanager
def temp_skip(self):
prev_status = self.skip_animations
self.skip_animations = True
try:
yield
finally:
if not prev_status:
self.stop_skipping()
@contextmanager
def temp_progress_bar(self):
prev_progress = self.show_animation_progress
self.show_animation_progress = progress_bar
self.show_animation_progress = True
try:
yield
finally:
self.show_animation_progress = prev_progress
if record:
self.camera.use_window_fbo(False)
self.file_writer.begin_insert()
self.checkpoint_manager.checkpoint_paste(self)
if record:
@contextmanager
def temp_record(self):
self.camera.use_window_fbo(False)
self.file_writer.begin_insert()
try:
yield
finally:
self.file_writer.end_insert()
self.camera.use_window_fbo(True)
self.stop_skipping()
self.show_animation_progress = prev_progress
def clear_checkpoints(self):
self.checkpoint_manager.clear_checkpoints()
def temp_config_change(self, skip=False, record=False, progress_bar=False):
stack = ExitStack()
if skip:
stack.enter_context(self.temp_skip())
if record:
stack.enter_context(self.temp_record())
if progress_bar:
stack.enter_context(self.temp_progress_bar())
return stack
def is_window_closing(self):
return self.window and (self.window.is_closing or self.quit_interaction)

View file

@ -1,7 +1,9 @@
from __future__ import annotations
import inspect
import pyperclip
import traceback
from IPython.core.getipython import get_ipython
from IPython.terminal import pt_inputhooks
from IPython.terminal.embed import InteractiveShellEmbed
@ -13,137 +15,144 @@ 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()
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from manimlib.scene.scene import Scene
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
class InteractiveSceneEmbed:
def __init__(self, scene: Scene):
self.scene = scene
self.checkpoint_manager = CheckpointManager()
# 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))
exception_mode = manim_config.embed.exception_mode
self.shell = self.get_ipython_shell_for_embedded_scene()
self.enable_gui()
self.ensure_frame_update_post_cell()
self.ensure_flash_on_error()
return InteractiveShellEmbed(
user_module=module,
display_banner=False,
xmode=exception_mode
)
def launch(self):
self.shell()
def get_ipython_shell_for_embedded_scene(self) -> InteractiveShellEmbed:
"""
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
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,
undo=scene.undo,
redo=scene.redo,
i2g=scene.i2g,
i2m=scene.i2m,
checkpoint_paste=scene.checkpoint_paste,
reload=reload_scene # Defined below
)
# 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(self.get_shortcuts())
exception_mode = manim_config.embed.exception_mode
return InteractiveShellEmbed(
user_module=module,
display_banner=False,
xmode=exception_mode
)
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()
def get_shortcuts(self):
"""
A few custom shortcuts useful to have in the interactive shell namespace
"""
scene = self.scene
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,
undo=scene.undo,
redo=scene.redo,
i2g=scene.i2g,
i2m=scene.i2m,
checkpoint_paste=self.checkpoint_paste,
clear_checkpoints=self.checkpoint_manager.clear_checkpoints,
reload=self.reload_scene # Defined below
)
pt_inputhooks.register("manim", inputhook)
shell.enable_gui("manim")
def enable_gui(self):
"""Enables gui interactions during the embed"""
def inputhook(context):
while not context.input_is_ready():
if not self.scene.is_window_closing():
self.scene.update_frame(dt=0)
if self.scene.is_window_closing():
self.shell.ask_exit()
pt_inputhooks.register("manim", inputhook)
self.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)
def ensure_frame_update_post_cell(self):
"""Ensure the scene updates its frame after each ipython cell"""
def post_cell_func(*args, **kwargs):
if not self.scene.is_window_closing():
self.scene.update_frame(dt=0, force_draw=True)
shell.events.register("post_run_cell", post_cell_func)
self.shell.events.register("post_run_cell", post_cell_func)
def ensure_flash_on_error(self):
"""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
print(''.join(traceback.format_exception(etype, evalue, tb)))
rect = FullScreenRectangle().set_stroke(RED, 30).set_fill(opacity=0)
rect.fix_in_frame()
self.scene.play(VFadeInThenOut(rect, run_time=0.5))
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)
rect = FullScreenRectangle().set_stroke(RED, 30).set_fill(opacity=0)
rect.fix_in_frame()
scene.play(VFadeInThenOut(rect, run_time=0.5))
self.shell.set_custom_exc((Exception,), custom_exc)
shell.set_custom_exc((Exception,), custom_exc)
def reload_scene(self, embed_line: int | None = None) -> None:
"""
Reloads the scene just like the `manimgl` command would do with the
same arguments that were provided for the initial startup. This allows
for quick iteration during scene development since we don't have to exit
the IPython kernel and re-run the `manimgl` command again. The GUI stays
open during the reload.
If `embed_line` is provided, the scene will be reloaded at that line
number. This corresponds to the `linemarker` param of the
`extract_scene.insert_embed_line_to_module()` method.
def reload_scene(embed_line: int | None = None) -> None:
"""
Reloads the scene just like the `manimgl` command would do with the
same arguments that were provided for the initial startup. This allows
for quick iteration during scene development since we don't have to exit
the IPython kernel and re-run the `manimgl` command again. The GUI stays
open during the reload.
Before reload, the scene is cleared and the entire state is reset, such
that we can start from a clean slate. This is taken care of by the
run_scenes function in __main__.py, which will catch the error raised by the
`exit_raise` magic command that we invoke here.
If `embed_line` is provided, the scene will be reloaded at that line
number. This corresponds to the `linemarker` param of the
`extract_scene.insert_embed_line_to_module()` method.
Note that we cannot define a custom exception class for this error,
since the IPython kernel will swallow any exception. While we can catch
such an exception in our custom exception handler registered with the
`set_custom_exc` method, we cannot break out of the IPython shell by
this means.
"""
# Update the global run configuration.
run_config = manim_config.run
run_config.is_reload = True
if embed_line:
run_config.embed_line = embed_line
Before reload, the scene is cleared and the entire state is reset, such
that we can start from a clean slate. This is taken care of by the
run_scenes function in __main__.py, which will catch the error raised by the
`exit_raise` magic command that we invoke here.
print("Reloading...")
self.shell.run_line_magic("exit_raise", "")
Note that we cannot define a custom exception class for this error,
since the IPython kernel will swallow any exception. While we can catch
such an exception in our custom exception handler registered with the
`set_custom_exc` method, we cannot break out of the IPython shell by
this means.
"""
shell = get_ipython()
if not shell:
return
# Update the global run configuration.
run_config = manim_config.run
run_config.is_reload = True
if embed_line:
run_config.embed_line = embed_line
print("Reloading...")
shell.run_line_magic("exit_raise", "")
def checkpoint_paste(
self,
skip: bool = False,
record: bool = False,
progress_bar: bool = True
):
with self.scene.temp_config_change(skip, record, progress_bar):
self.checkpoint_manager.checkpoint_paste(self.shell, self.scene)
class CheckpointManager:
checkpoint_states: dict[str, list[tuple[Mobject, Mobject]]] = dict()
def checkpoint_paste(self, scene):
def checkpoint_paste(self, shell, scene):
"""
Used during interactive development to run (or re-run)
a block of scene code.
@ -152,25 +161,20 @@ class CheckpointManager:
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):
def get_leading_comment(code_string: str) -> str:
leading_line = code_string.partition("\n")[0].lstrip()
if leading_line.startswith("#"):
return leading_line
return None
return ""
def handle_checkpoint_key(self, scene, key: str):
if key is None:
if not key:
return
elif key in self.checkpoint_states:
# Revert to checkpoint