From 636fb3a45b4726ca308043a665b60f4e9efc5e2d Mon Sep 17 00:00:00 2001 From: Grant Sanderson Date: Mon, 9 Dec 2024 13:53:03 -0600 Subject: [PATCH] Factor interactive embed logic out of Scene class --- manimlib/scene/scene.py | 134 ++----------------------- manimlib/scene/scene_embed.py | 178 ++++++++++++++++++++++++++++++++++ 2 files changed, 184 insertions(+), 128 deletions(-) create mode 100644 manimlib/scene/scene_embed.py diff --git a/manimlib/scene/scene.py b/manimlib/scene/scene.py index d35aa7c2..df1a4a38 100644 --- a/manimlib/scene/scene.py +++ b/manimlib/scene/scene.py @@ -1,41 +1,34 @@ from __future__ import annotations from collections import OrderedDict -import inspect import os import platform -import pyperclip import random import time -import re from functools import wraps -from IPython.terminal import pt_inputhooks -from IPython.terminal.embed import InteractiveShellEmbed from pyglet.window import key as PygletWindowKeys import numpy as np from tqdm.auto import tqdm as ProgressDisplay from manimlib.animation.animation import prepare_animation -from manimlib.animation.fading import VFadeInThenOut from manimlib.camera.camera import Camera from manimlib.camera.camera_frame import CameraFrame -from manimlib.module_loader import ModuleLoader from manimlib.constants import ARROW_SYMBOLS from manimlib.constants import DEFAULT_WAIT_TIME -from manimlib.constants import RED from manimlib.event_handler import EVENT_DISPATCHER from manimlib.event_handler.event_type import EventType from manimlib.logger import log from manimlib.reload_manager import reload_manager -from manimlib.mobject.frame import FullScreenRectangle from manimlib.mobject.mobject import _AnimationBuilder from manimlib.mobject.mobject import Group 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 CheckpointManager from manimlib.scene.scene_file_writer import SceneFileWriter from manimlib.utils.family_ops import extract_mobject_family_members from manimlib.utils.family_ops import recursive_mobject_remove @@ -126,6 +119,7 @@ class Scene(object): 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 = [] @@ -210,80 +204,9 @@ class Scene(object): close_scene_on_exit: bool = True, show_animation_progress: bool = False, ) -> None: - if not self.window: - # Embed is only relevant for interactive development with a Window - return - self.stop_skipping() - self.update_frame(force_draw=True) - self.save_state() self.show_animation_progress = show_animation_progress - # Create embedded IPython terminal configured to have access to - # the local namespace of the caller - caller_frame = inspect.currentframe().f_back - module = ModuleLoader.get_module(caller_frame.f_globals["__file__"]) - shell = InteractiveShellEmbed( - user_module=module, - display_banner=False, - xmode=self.embed_exception_mode - ) - self.shell = shell - - # Add a few custom shortcuts to that local namespace - local_ns = dict(caller_frame.f_locals) - local_ns.update( - play=self.play, - wait=self.wait, - add=self.add, - remove=self.remove, - clear=self.clear, - focus=self.focus, - save_state=self.save_state, - reload=self.reload, - undo=self.undo, - redo=self.redo, - i2g=self.i2g, - i2m=self.i2m, - checkpoint_paste=self.checkpoint_paste, - touch=lambda: shell.enable_gui("manim"), - notouch=lambda: shell.enable_gui(None), - ) - - # Update the shell module with the caller's locals + shortcuts - module.__dict__.update(local_ns) - - # Enables gui interactions during the embed - def inputhook(context): - while not context.input_is_ready(): - if not self.is_window_closing(): - self.update_frame(dt=0) - if self.is_window_closing(): - shell.ask_exit() - - pt_inputhooks.register("manim", inputhook) - shell.enable_gui("manim") - - # Operation to run after each ipython command - def post_cell_func(*args, **kwargs): - if not self.is_window_closing(): - self.update_frame(dt=0, force_draw=True) - - shell.events.register("post_run_cell", post_cell_func) - - # 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 self.embed_error_sound: - os.system("printf '\a'") - rect = FullScreenRectangle().set_stroke(RED, 30).set_fill(opacity=0) - rect.fix_in_frame() - self.play(VFadeInThenOut(rect, run_time=0.5)) - - shell.set_custom_exc((Exception,), custom_exc) - - # Launch shell - shell() + interactive_scene_embed(self) # End scene when exiting an embed if close_scene_on_exit: @@ -760,37 +683,6 @@ class Scene(object): revert to the state of the scene the first time this function was called on a block of code starting with that comment. """ - if self.shell is None or self.window is None: - raise Exception( - "Scene.checkpoint_paste cannot be called outside of " + - "an ipython shell" - ) - - pasted = pyperclip.paste() - lines = pasted.split("\n") - - # Commented lines trigger saved checkpoints - if lines[0].lstrip().startswith("#"): - if lines[0] not in self.checkpoint_states: - self.checkpoint(lines[0]) - else: - self.revert_to_checkpoint(lines[0]) - - # Copied methods of a scene are handled specially - # A bit hacky, yes, but convenient - method_pattern = r"^def\s+([a-zA-Z_]\w*)\s*\(self.*\):" - method_names = re.findall(method_pattern ,lines[0].strip()) - if method_names: - method_name = method_names[0] - indent = " " * lines[0].index(lines[0].strip()) - pasted = "\n".join([ - # Remove self from function signature - re.sub(r"self(,\s*)?", "", lines[0]), - *lines[1:], - # Attach to scene via self.func_name = func_name - f"{indent}self.{method_name} = {method_name}" - ]) - # Keep track of skipping and progress bar status self.skip_animations = skip @@ -801,7 +693,7 @@ class Scene(object): self.camera.use_window_fbo(False) self.file_writer.begin_insert() - self.shell.run_cell(pasted) + self.checkpoint_manager.checkpoint_paste(self) if record: self.file_writer.end_insert() @@ -810,22 +702,8 @@ class Scene(object): self.stop_skipping() self.show_animation_progress = prev_progress - def checkpoint(self, key: str): - self.checkpoint_states[key] = self.get_state() - - def revert_to_checkpoint(self, key: str): - if key not in self.checkpoint_states: - log.error(f"No checkpoint at {key}") - return - 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) - - self.restore_state(self.checkpoint_states[key]) - def clear_checkpoints(self): - self.checkpoint_states = dict() + self.checkpoint_manager.clear_checkpoints() def save_mobject_to_file(self, mobject: Mobject, file_path: str | None = None) -> None: if file_path is None: diff --git a/manimlib/scene/scene_embed.py b/manimlib/scene/scene_embed.py new file mode 100644 index 00000000..43f4be5f --- /dev/null +++ b/manimlib/scene/scene_embed.py @@ -0,0 +1,178 @@ +import inspect +import pyperclip +import re + +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): + if not scene.window: + # Embed is only relevant for interactive development with a Window + return + scene.stop_skipping() + scene.update_frame(force_draw=True) + scene.save_state() + + shell = get_ipython_shell_for_embedded_scene(scene) + scene.shell = shell # It would be better not to add attributes to scene here + 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 + module = ModuleLoader.get_module(caller_frame.f_globals["__file__"]) + shell = InteractiveShellEmbed( + user_module=module, + display_banner=False, + xmode=scene.embed_exception_mode + ) + + # Update the module's namespace to match include local variables + module.__dict__.update(caller_frame.f_locals) + module.__dict__.update(get_shortcuts(scene)) + + return shell + + +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) + code_string = self.handle_method_definitions(code_string) + + shell.run_cell(code_string) + + @staticmethod + def handle_method_definitions(code_string: str): + lines = code_string.split("\n") + # Copied methods of a scene are handled specially + # A bit hacky, yes, but convenient + method_pattern = r"^def\s+([a-zA-Z_]\w*)\s*\(self.*\):" + method_names = re.findall(method_pattern, lines[0].strip()) + if method_names: + method_name = method_names[0] + indent = " " * lines[0].index(lines[0].strip()) + return "\n".join([ + # Remove self from function signature + re.sub(r"self(,\s*)?", "", lines[0]), + *lines[1:], + # Attach to scene via self.func_name = func_name + f"{indent}self.{method_name} = {method_name}" + ]) + return 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()