๐Ÿ“ฆ 3b1b / manim

๐Ÿ“„ scene_embed.py ยท 153 lines
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153import inspect
import pyperclip
import os

from IPython.core.getipython import get_ipython
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):
    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()


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

    # 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))

    return InteractiveShellEmbed(
        user_module=module,
        display_banner=False,
        xmode=scene.embed_exception_mode
    )


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)
        shell.run_cell(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()