diff --git a/manimlib/config.py b/manimlib/config.py index cd6dd4a7..b4420a92 100644 --- a/manimlib/config.py +++ b/manimlib/config.py @@ -9,6 +9,7 @@ import os import screeninfo import sys import yaml +from pathlib import Path from functools import lru_cache @@ -144,12 +145,8 @@ def parse_cli(): ) parser.add_argument( "-e", "--embed", - nargs="?", - const="", help="Creates a new file where the line `self.embed` is inserted " + \ - "into the Scenes construct method. " + \ - "If a string is passed in, the line will be inserted below the " + \ - "last line of code including that string." + "at the corresponding line number" ) parser.add_argument( "-r", "--resolution", @@ -210,12 +207,19 @@ def get_manim_dir(): return os.path.abspath(os.path.join(manimlib_dir, "..")) -def get_indent(line: str): - return len(line) - len(line.lstrip()) +def get_indent(code_lines: list[str], line_number: int): + for line in code_lines[line_number:0:-1]: + if len(line.strip()) == 0: + continue + n_spaces = len(line) - len(line.lstrip()) + if line.endswith(":"): + n_spaces += 4 + return n_spaces + return 0 def get_module_with_inserted_embed_line( - file_name: str, scene_name: str, line_marker: str + file_name: str, scene_name: str, line_number: int ): """ This is hacky, but convenient. When user includes the argument "-e", it will try @@ -223,59 +227,24 @@ def get_module_with_inserted_embed_line( construct method. If there is an argument passed in, it will insert the line after the last line in the sourcefile which includes that string. """ - with open(file_name, 'r') as fp: - lines = fp.readlines() + lines = Path(file_name).read_text().splitlines() - try: - scene_line_number = next( - i for i, line in enumerate(lines) - if line.startswith(f"class {scene_name}") - ) - except StopIteration: + scene_line_numbers = [ + n for n, line in enumerate(lines) + if line.startswith("class SurfaceTest") + ] + if len(scene_line_numbers) == 0: log.error(f"No scene {scene_name}") return + scene_line_number = scene_line_numbers[0] - prev_line_num = -1 - n_spaces = None - if len(line_marker) == 0: - # Find the end of the construct method - in_construct = False - for index in range(scene_line_number, len(lines) - 1): - line = lines[index] - if line.lstrip().startswith("def construct"): - in_construct = True - n_spaces = get_indent(line) + 4 - elif in_construct: - if len(line.strip()) > 0 and get_indent(line) < (n_spaces or 0): - prev_line_num = index - 1 - break - if prev_line_num < 0: - prev_line_num = len(lines) - 1 - elif line_marker.isdigit(): - # Treat the argument as a line number - prev_line_num = int(line_marker) - 1 - elif len(line_marker) > 0: - # Treat the argument as a string - try: - prev_line_num = next( - i - for i in range(scene_line_number, len(lines) - 1) - if line_marker in lines[i] - ) - except StopIteration: - log.error(f"No lines matching {line_marker}") - sys.exit(2) - - # Insert the embed line, rewrite file, then write it back when done - if n_spaces is None: - n_spaces = get_indent(lines[prev_line_num]) - inserted_line = " " * n_spaces + "self.embed()\n" + n_spaces = get_indent(lines, line_number - 1) + inserted_line = " " * n_spaces + "self.embed()" new_lines = list(lines) - new_lines.insert(prev_line_num + 1, inserted_line) + new_lines.insert(line_number, inserted_line) new_file = file_name.replace(".py", "_insert_embed.py") - with open(new_file, 'w') as fp: - fp.writelines(new_lines) + Path(new_file).write_text("\n".join(new_lines)) from manimlib.reload_manager import reload_manager module = ModuleLoader.get_module(new_file, is_during_reload=reload_manager.is_reload) @@ -293,7 +262,7 @@ def get_scene_module(args: Namespace) -> Module: return ModuleLoader.get_module(args.file) else: return get_module_with_inserted_embed_line( - args.file, args.scene_names[0], args.embed + args.file, args.scene_names[0], int(args.embed) )