build

#!/bin/python3

"""A script for building and renaming all of the Manim scenes."""
import subprocess
import shutil
import os
import argparse
import re

from glob import glob


VIDEO_DIRECTORY = "video"

parser = argparse.ArgumentParser()

parser.add_argument(
    "-s", "--scenes",
    nargs='*',
    help="the name of the scenes to build; builds all when omitted",
    default=None,
)

parser.add_argument(
    "-q", "--quality",
    choices=["m", "h", "k"],
    default="k",
    help="the quality of the video (4K by default)"
)

parser.add_argument(
    "-r", "--rename-only",
    help="only perform scene renaming",
    action='store_true',
)

os.chdir(os.path.dirname(os.path.abspath(__file__)))

arguments = parser.parse_args()


def render_scene(scene: str):
    """Render a single scene."""
    quality_mapping = {
        "m": (1280, 720, 30),
        "h": (1920, 1080, 60),
        "k": (3840, 2160, 60),
    }

    w, h, f = quality_mapping[arguments.quality]

    build_dir = os.path.abspath(os.path.dirname(__file__))

    if os.path.isfile(os.path.join(build_dir, ".short")):
        w, h = h, w

    args = ["python3", "-m", "manim", "scenes.py", "--fps", str(f), "-r", f"{w},{h}", "--disable_caching", scene]

    if scene.lower().startswith("transparent"):
        args.append("-t")

    process = subprocess.Popen(args)
    process.communicate()

    if process.returncode != 0:
        print(f"\nBuild failed with exit code {process.returncode}")
        quit()


def rename_scene(scene: str):
    """Use Manim's partial movie txt file to rename the videos to 1.mp4, 2.mp4 etc..."""
    partial_file_path = os.path.join(VIDEO_DIRECTORY, scene, "partial_movie_file_list.txt")

    if os.path.exists(partial_file_path):
        with open(partial_file_path) as f:
            for i, video in enumerate(f.read().splitlines()[1:]):
                path, name = os.path.split(video[11:-1])

                ext = os.path.splitext(name)[1][1:]

                original_path = os.path.join(path, name)
                changed_path = os.path.join(path, f"{i+1}.{ext}")

                os.rename(original_path, changed_path)

        os.remove(partial_file_path)
    else:
        print(f"WARNING: Partial movie file list for scene '{scene}' doesn't exists, not doing anything!")


scenes = []

# please lord forgive me for I have sinned
# this is needed because classes ending with Transparent are going to be rendered transparently
# and classes that contain "Test" are ignored
with open("scenes.py", "r") as f:
    for line in f.read().splitlines():
        if match := re.match(r"\s*class\s+(.+?)\(.*Scene\)\s*:", line):
            scenes.append(match.group(1))

if arguments.scenes is None or len(arguments.scenes) == 0:
    # remove everything when no scene is specified
    if not arguments.rename_only:
        if os.path.exists(VIDEO_DIRECTORY):
            shutil.rmtree(VIDEO_DIRECTORY)

    for scene in scenes:
        # skip test scenes
        if "Test" in scene:
            continue

        if not arguments.rename_only:
            render_scene(scene)

        rename_scene(scene)
else:
    for scene in arguments.scenes:
        scene_folder = os.path.join(VIDEO_DIRECTORY, scene)

        # remove only the scene and its folders when the scene is specified
        if not arguments.rename_only:
            if os.path.exists(scene_folder):
                shutil.rmtree(scene_folder)

            for ext in ["mp4", "mov"]:
                scene_video = os.path.join(VIDEO_DIRECTORY, scene + f".{ext}")
                if os.path.exists(scene_video):
                    os.remove(scene_video)

        if not arguments.rename_only:
            render_scene(scene)

        rename_scene(scene)