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(
-        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:
+    @contextmanager
+    def temp_record(self):
         self.camera.use_window_fbo(False)
         self.file_writer.begin_insert()
-
-        self.checkpoint_manager.checkpoint_paste(self)
-
-        if record:
+        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)

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,20 +15,25 @@ 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):
+class InteractiveSceneEmbed:
+    def __init__(self, scene: Scene):
+        self.scene = scene
+        self.checkpoint_manager = CheckpointManager()
+
+        self.shell = self.get_ipython_shell_for_embedded_scene()
+        self.enable_gui()
+        self.ensure_frame_update_post_cell()
+        self.ensure_flash_on_error()
+
+    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
@@ -38,7 +45,7 @@ def get_ipython_shell_for_embedded_scene(scene):
         # 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))
+        module.__dict__.update(self.get_shortcuts())
         exception_mode = manim_config.embed.exception_mode
 
         return InteractiveShellEmbed(
@@ -47,11 +54,11 @@ def get_ipython_shell_for_embedded_scene(scene):
             xmode=exception_mode
         )
 
-
-def get_shortcuts(scene):
+    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,
@@ -64,46 +71,43 @@ def get_shortcuts(scene):
             redo=scene.redo,
             i2g=scene.i2g,
             i2m=scene.i2m,
-        checkpoint_paste=scene.checkpoint_paste,
-        reload=reload_scene  # Defined below
+            checkpoint_paste=self.checkpoint_paste,
+            clear_checkpoints=self.checkpoint_manager.clear_checkpoints,
+            reload=self.reload_scene  # Defined below
         )
 
-
-def enable_gui(shell, scene):
+    def enable_gui(self):
         """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()
+                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)
-    shell.enable_gui("manim")
+        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"""
         def post_cell_func(*args, **kwargs):
-        if not scene.is_window_closing():
-            scene.update_frame(dt=0, force_draw=True)
+            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(shell, scene):
+    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
-        shell.showtraceback((etype, evalue, tb), tb_offset=tb_offset)
+            print(''.join(traceback.format_exception(etype, evalue, tb)))
             rect = FullScreenRectangle().set_stroke(RED, 30).set_fill(opacity=0)
             rect.fix_in_frame()
-        scene.play(VFadeInThenOut(rect, run_time=0.5))
+            self.scene.play(VFadeInThenOut(rect, run_time=0.5))
 
-    shell.set_custom_exc((Exception,), custom_exc)
+        self.shell.set_custom_exc((Exception,), custom_exc)
 
-
-def reload_scene(embed_line: int | None = None) -> None:
+    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
@@ -126,10 +130,6 @@ def reload_scene(embed_line: int | None = None) -> None:
         `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
@@ -137,13 +137,22 @@ def reload_scene(embed_line: int | None = None) -> None:
             run_config.embed_line = embed_line
 
         print("Reloading...")
-    shell.run_line_magic("exit_raise", "")
+        self.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