3b1b-manim/manimlib/extract_scene.py

import importlib
import inspect
import itertools as it
import os
import subprocess as sp
import sys
import traceback

from manimlib.scene.scene import Scene
from manimlib.utils.sounds import play_error_sound
from manimlib.utils.sounds import play_finish_sound
import manimlib.constants


def handle_scene(scene, **config):
    import platform
    if config["quiet"]:
        curr_stdout = sys.stdout
        sys.stdout = open(os.devnull, "w")

    if config["show_last_frame"]:
        scene.save_image(mode=config["saved_image_mode"])
    open_file = any([
        config["show_last_frame"],
        config["open_video_upon_completion"],
        config["show_file_in_finder"]
    ])
    if open_file:
        commands = ["open"]
        if (platform.system() == "Linux"):
            commands = ["xdg-open"]
        elif (platform.system() == "Windows"):
            commands = ["start"]

        if config["show_file_in_finder"]:
            commands.append("-R")

        if config["show_last_frame"]:
            commands.append(scene.get_image_file_path())
        else:
            commands.append(scene.get_movie_file_path())
        # commands.append("-g")
        FNULL = open(os.devnull, 'w')
        sp.call(commands, stdout=FNULL, stderr=sp.STDOUT)
        FNULL.close()

    if config["quiet"]:
        sys.stdout.close()
        sys.stdout = curr_stdout


def is_child_scene(obj, module):
    if not inspect.isclass(obj):
        return False
    if not issubclass(obj, Scene):
        return False
    if obj == Scene:
        return False
    if not obj.__module__.startswith(module.__name__):
        return False
    return True


def prompt_user_for_choice(name_to_obj):
    num_to_name = {}
    names = sorted(name_to_obj.keys())
    for count, name in zip(it.count(1), names):
        print("%d: %s" % (count, name))
        num_to_name[count] = name
    try:
        user_input = input(manimlib.constants.CHOOSE_NUMBER_MESSAGE)
        return [
            name_to_obj[num_to_name[int(num_str)]]
            for num_str in user_input.split(",")
        ]
    except KeyError:
        print(manimlib.constants.INVALID_NUMBER_MESSAGE)
        sys.exit(2)
        user_input = input(manimlib.constants.CHOOSE_NUMBER_MESSAGE)
        return [
            name_to_obj[num_to_name[int(num_str)]]
            for num_str in user_input.split(",")
        ]
    except EOFError:
        sys.exit(1)


def get_scene_classes(scene_names_to_classes, config):
    if len(scene_names_to_classes) == 0:
        print(manimlib.constants.NO_SCENE_MESSAGE)
        return []
    if config["scene_name"] in scene_names_to_classes:
        return [scene_names_to_classes[config["scene_name"]]]
    if config["scene_name"] != "":
        print(manimlib.constants.SCENE_NOT_FOUND_MESSAGE)
        return []
    if config["write_all"]:
        return list(scene_names_to_classes.values())
    return prompt_user_for_choice(scene_names_to_classes)


def main(config):
    module = config["module"]
    scene_names_to_classes = dict(
        inspect.getmembers(module, lambda x: is_child_scene(x, module)))

    scene_kwargs = dict([
        (key, config[key])
        for key in [
            "camera_config",
            "frame_duration",
            "skip_animations",
            "write_to_movie",
            "save_pngs",
            "movie_file_extension",
            "start_at_animation_number",
            "end_at_animation_number",
        ]
    ])

    scene_kwargs["name"] = config["output_name"]
    if config["save_pngs"]:
        print("We are going to save a PNG sequence as well...")
        scene_kwargs["save_pngs"] = True
        scene_kwargs["pngs_mode"] = config["saved_image_mode"]

    for SceneClass in get_scene_classes(scene_names_to_classes, config):
        try:
            handle_scene(SceneClass(**scene_kwargs), **config)
            play_finish_sound()
            sys.exit(0)
        except Exception:
            print("\n\n")
            traceback.print_exc()
            print("\n\n")
            play_error_sound()
            sys.exit(2)


if __name__ == "__main__":
    main()