Update InteractiveSceneEmbed (#2267)

* 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

manimlib/module_loader.py

@@ -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__"):

manimlib/scene/scene.py

@@ -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(
+    @contextmanager
-        self,
+    def temp_skip(self):
-        skip: bool = False,
+        prev_status = self.skip_animations
-        record: bool = False,
+        self.skip_animations = True
-        progress_bar: bool = True
+        try:
-    ):
+            yield
-        """
+        finally:
-        Used during interactive development to run (or re-run)
+            if not prev_status:
-        a block of scene code.
+                self.stop_skipping()
-
-        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_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:
+    @contextmanager
-            self.camera.use_window_fbo(False)
+    def temp_record(self):
-            self.file_writer.begin_insert()
+        self.camera.use_window_fbo(False)
-
+        self.file_writer.begin_insert()
-        self.checkpoint_manager.checkpoint_paste(self)
+        try:
-
+            yield
-        if record:
+        finally:
             self.file_writer.end_insert()
             self.camera.use_window_fbo(True)
 
-        self.stop_skipping()
+    def temp_config_change(self, skip=False, record=False, progress_bar=False):
-        self.show_animation_progress = prev_progress
+        stack = ExitStack()
-
+        if skip:
-    def clear_checkpoints(self):
+            stack.enter_context(self.temp_skip())
-        self.checkpoint_manager.clear_checkpoints()
+        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)

manimlib/scene/scene_embed.py

@@ -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):
+from typing import TYPE_CHECKING
-    scene.stop_skipping()
+if TYPE_CHECKING:
-    scene.update_frame(force_draw=True)
+    from manimlib.scene.scene import Scene
-
-    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):
+class InteractiveSceneEmbed:
-    """
+    def __init__(self, scene: Scene):
-    Create embedded IPython terminal configured to have access to
+        self.scene = scene
-    the local namespace of the caller
+        self.checkpoint_manager = CheckpointManager()
-    """
-    # 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
+        self.shell = self.get_ipython_shell_for_embedded_scene()
-    module = ModuleLoader.get_module(caller_frame.f_globals["__file__"])
+        self.enable_gui()
-    module.__dict__.update(caller_frame.f_locals)
+        self.ensure_frame_update_post_cell()
-    module.__dict__.update(get_shortcuts(scene))
+        self.ensure_flash_on_error()
-    exception_mode = manim_config.embed.exception_mode
 
-    return InteractiveShellEmbed(
+    def launch(self):
-        user_module=module,
+        self.shell()
-        display_banner=False,
-        xmode=exception_mode
-    )
 
+    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):
+        # Update the module's namespace to include local variables
-    """
+        module = ModuleLoader.get_module(caller_frame.f_globals["__file__"])
-    A few custom shortcuts useful to have in the interactive shell namespace
+        module.__dict__.update(caller_frame.f_locals)
-    """
+        module.__dict__.update(self.get_shortcuts())
-    return dict(
+        exception_mode = manim_config.embed.exception_mode
-        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
-    )
 
+        return InteractiveShellEmbed(
+            user_module=module,
+            display_banner=False,
+            xmode=exception_mode
+        )
 
-def enable_gui(shell, scene):
+    def get_shortcuts(self):
-    """Enables gui interactions during the embed"""
+        """
-    def inputhook(context):
+        A few custom shortcuts useful to have in the interactive shell namespace
-        while not context.input_is_ready():
+        """
-            if not scene.is_window_closing():
+        scene = self.scene
-                scene.update_frame(dt=0)
+        return dict(
-        if scene.is_window_closing():
+            play=scene.play,
-            shell.ask_exit()
+            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)
+    def enable_gui(self):
-    shell.enable_gui("manim")
+        """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):
+    def ensure_frame_update_post_cell(self):
-    """Ensure the scene updates its frame after each ipython cell"""
+        """Ensure the scene updates its frame after each ipython cell"""
-    def post_cell_func(*args, **kwargs):
+        def post_cell_func(*args, **kwargs):
-        if not scene.is_window_closing():
+            if not self.scene.is_window_closing():
-            scene.update_frame(dt=0, force_draw=True)
+                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):
+        self.shell.set_custom_exc((Exception,), custom_exc)
-    """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))
 
-    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:
+        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
-    Reloads the scene just like the `manimgl` command would do with the
+        run_scenes function in __main__.py, which will catch the error raised by the
-    same arguments that were provided for the initial startup. This allows
+        `exit_raise` magic command that we invoke here.
-    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
+        Note that we cannot define a custom exception class for this error,
-    number. This corresponds to the `linemarker` param of the
+        since the IPython kernel will swallow any exception. While we can catch
-    `extract_scene.insert_embed_line_to_module()` method.
+        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
+        print("Reloading...")
-    that we can start from a clean slate. This is taken care of by the
+        self.shell.run_line_magic("exit_raise", "")
-    run_scenes function in __main__.py, which will catch the error raised by the
-    `exit_raise` magic command that we invoke here.
 
-    Note that we cannot define a custom exception class for this error,
+    def checkpoint_paste(
-    since the IPython kernel will swallow any exception. While we can catch
+        self,
-    such an exception in our custom exception handler registered with the
+        skip: bool = False,
-    `set_custom_exc` method, we cannot break out of the IPython shell by
+        record: bool = False,
-    this means.
+        progress_bar: bool = True
-    """
+    ):
-    shell = get_ipython()
+        with self.scene.temp_config_change(skip, record, progress_bar):
-    if not shell:
+            self.checkpoint_manager.checkpoint_paste(self.shell, self.scene)
-        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", "")
 
 
 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