Add reload() command for interactive scene reloading (#2240)

* Init reload command (lots of things not working yet)

* Add back in class line (accidentally deleted)

* Add back in key modifiers (accidentally deleted)

* Unpack tuple from changed `get_module`

* Init MainRunManager & respawn IPython shell

* Init cleanup of scenes from manager

* Restore string quotes

* Still take `self.preview` into account

* Remove left-over code from module experimentation

* Remove double window activation

* Reset scenes array in RunManager

* Move self.args None check up

* Use first available window

* Don't use constructor for RunManager

* Use self. syntax

* Init moderngl context manually

* Add some comments for failed attempts to reset scene

* Reuse existing shell (this fixed the bug 🎉)

* Remove unused code

* Remove unnecessary intermediate ReloadSceneException

* Allow users to finally exit

* Rename main_run_manager to reload_manager

* Add docstrings to `ReloadManager`

* Improve reset management in window

* Clarify why we use magic exit_raise command

* Add comment about window reuse

* Improve docstrings in ReloadManager & handle case of 0 scenes

* Set scene and title earlier

* Run linter suggestions

manimlib/__main__.py

@@ -1,12 +1,15 @@
 #!/usr/bin/env python
 from manimlib import __version__
 import manimlib.config
-import manimlib.extract_scene
 import manimlib.logger
 import manimlib.utils.init_config
+from manimlib.reload_manager import reload_manager
 
 
 def main():
+    """
+    Main entry point for ManimGL.
+    """
     print(f"ManimGL \033[32mv{__version__}\033[0m")
 
     args = manimlib.config.parse_cli()
@@ -17,12 +20,10 @@ def main():
 
     if args.config:
         manimlib.utils.init_config.init_customization()
-    else:
-        config = manimlib.config.get_configuration(args)
-        scenes = manimlib.extract_scene.main(config)
+        return
 
-        for scene in scenes:
-            scene.run()
+    reload_manager.args = args
+    reload_manager.run()
 
 
 if __name__ == "__main__":

manimlib/reload_manager.py

@@ -0,0 +1,84 @@
+from typing import Any
+from IPython.terminal.embed import KillEmbedded
+
+
+class ReloadManager:
+    """
+    Manages the loading and running of scenes and is called directly from the
+    main entry point of ManimGL.
+
+    The name "reload" comes from the fact that this class handles the
+    reinitialization of scenes when requested by the user via the `reload()`
+    command in the IPython shell.
+    """
+
+    args: Any = None
+    scenes: list[Any] = []
+    window = None
+
+    # The line number to load the scene from when reloading
+    start_at_line = None
+
+    def set_new_start_at_line(self, start_at_line):
+        """
+        Sets/Updates the line number to load the scene from when reloading.
+        """
+        self.start_at_line = start_at_line
+
+    def run(self):
+        """
+        Runs the scenes in a loop and detects when a scene reload is requested.
+        """
+        while True:
+            try:
+                # blocking call since a scene will init an IPython shell()
+                self.retrieve_scenes_and_run(self.start_at_line)
+                return
+            except KillEmbedded:
+                # Requested via the `exit_raise` IPython runline magic
+                # by means of our scene.reload() command
+                print("Reloading...")
+
+                for scene in self.scenes:
+                    scene.tear_down()
+
+                self.scenes = []
+
+            except KeyboardInterrupt:
+                break
+
+    def retrieve_scenes_and_run(self, overwrite_start_at_line: int | None = None):
+        """
+        Creates a new configuration based on the CLI args and runs the scenes.
+        """
+        import manimlib.config
+        import manimlib.extract_scene
+
+        # Args
+        if self.args is None:
+            raise RuntimeError("Fatal error: No args were passed to the ReloadManager")
+        if overwrite_start_at_line is not None:
+            self.args.embed = str(overwrite_start_at_line)
+
+        # Args to Config
+        config = manimlib.config.get_configuration(self.args)
+        if self.window:
+            config["existing_window"] = self.window  # see scene initialization
+
+        # Scenes
+        self.scenes = manimlib.extract_scene.main(config)
+        if len(self.scenes) == 0:
+            print("No scenes found to run")
+            return
+
+        # Find first available window
+        for scene in self.scenes:
+            if scene.window is not None:
+                self.window = scene.window
+                break
+
+        for scene in self.scenes:
+            scene.run()
+
+
+reload_manager = ReloadManager()

manimlib/scene/scene.py

@@ -12,7 +12,6 @@ from functools import wraps
 
 from IPython.terminal import pt_inputhooks
 from IPython.terminal.embed import InteractiveShellEmbed
-from IPython.core.getipython import get_ipython
 from pyglet.window import key as PygletWindowKeys
 
 import numpy as np
@@ -29,6 +28,7 @@ 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
@@ -40,6 +40,7 @@ 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
 from manimlib.utils.iterables import batch_by_property
+from manimlib.window import Window
 
 from typing import TYPE_CHECKING
 
@@ -88,6 +89,7 @@ class Scene(object):
         show_animation_progress: bool = False,
         embed_exception_mode: str = "",
         embed_error_sound: bool = False,
+        existing_window: Window | None = None,
     ):
         self.skip_animations = skip_animations
         self.always_update_mobjects = always_update_mobjects
@@ -106,12 +108,16 @@ class Scene(object):
             config["samples"] = self.samples
         self.file_writer_config = {**self.default_file_writer_config, **file_writer_config}
 
-        # Initialize window, if applicable
+        # Initialize window, if applicable (and reuse window if provided during
+        # reload by means of the ReloadManager)
         if self.preview:
-            from manimlib.window import Window
+            if existing_window:
+                self.window = existing_window
+                self.window.update_scene(self)
+            else:
                 self.window = Window(scene=self, **self.window_config)
-            self.camera_config["window"] = self.window
                 self.camera_config["fps"] = 30  # Where's that 30 from?
+            self.camera_config["window"] = self.window
         else:
             self.window = None
 
@@ -152,6 +158,9 @@ class Scene(object):
     def __str__(self) -> str:
         return self.__class__.__name__
 
+    def get_window(self) -> Window | None:
+        return self.window
+
     def run(self) -> None:
         self.virtual_animation_start_time: float = 0
         self.real_animation_start_time: float = time.time()
@@ -224,6 +233,7 @@ class Scene(object):
         caller_frame = inspect.currentframe().f_back
         module = get_module(caller_frame.f_globals["__file__"])
         shell = InteractiveShellEmbed(user_module=module)
+        self.shell = shell
 
         # Add a few custom shortcuts to that local namespace
         local_ns = dict(caller_frame.f_locals)
@@ -235,6 +245,7 @@ class Scene(object):
             clear=self.clear,
             focus=self.focus,
             save_state=self.save_state,
+            reload=self.reload,
             undo=self.undo,
             redo=self.redo,
             i2g=self.i2g,
@@ -758,8 +769,7 @@ 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.
         """
-        shell = get_ipython()
-        if shell is None or self.window is None:
+        if self.shell is None or self.window is None:
             raise Exception(
                 "Scene.checkpoint_paste cannot be called outside of " +
                 "an ipython shell"
@@ -800,7 +810,7 @@ class Scene(object):
             self.camera.use_window_fbo(False)
             self.file_writer.begin_insert()
 
-        shell.run_cell(pasted)
+        self.shell.run_cell(pasted)
 
         if record:
             self.file_writer.end_insert()
@@ -985,6 +995,31 @@ class Scene(object):
     def on_close(self) -> None:
         pass
 
+    def reload(self, start_at_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 `start_at_line` is provided, the scene will be reloaded at that line
+        number. This corresponds to the `linemarker` param of the
+        `config.get_module_with_inserted_embed_line()` method.
+
+        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
+        ReloadManager, 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,
+        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.
+        """
+        reload_manager.set_new_start_at_line(start_at_line)
+        self.shell.run_line_magic("exit_raise", "")
+
     def focus(self) -> None:
         """
         Puts focus on the ManimGL window.

manimlib/window.py

@@ -35,16 +35,29 @@ class Window(PygletWindow):
 
         self.default_size = size
         self.default_position = self.find_initial_position(size)
-        self.scene = scene
         self.pressed_keys = set()
-        self.title = str(scene)
         self.size = size
 
+        self.update_scene(scene)
+
+    def update_scene(self, scene: Scene):
+        """
+        Resets the state and updates the scene associated to this window.
+
+        This is necessary when we want to reuse an *existing* window after a
+        `scene.reload()` was requested, which will create new scene instances.
+        """
+        self.pressed_keys.clear()
         self._has_undrawn_event = True
 
-        mglw.activate_context(window=self)
+        self.scene = scene
+        self.title = str(scene)
+
+        self.init_mgl_context()
+
         self.timer = Timer()
         self.config = mglw.WindowConfig(ctx=self.ctx, wnd=self, timer=self.timer)
+        mglw.activate_context(window=self, ctx=self.ctx)
         self.timer.start()
 
         self.to_default_position()