Fix n flag (#423)

* fixes n flag

* added test for n flag

* Apply suggestions from code review

Co-authored-by: Aathish Sivasubrahmanian <[email protected]>

Co-authored-by: Hugues Devimeux <[email protected]>
Co-authored-by: Aathish Sivasubrahmanian <[email protected]>

Author: 3 people

manim/config/config_utils.py

@@ -153,9 +153,7 @@ def _parse_file_writer_config(config_parser, args):
     }
 
     # For internal use (no CLI flag)
-    fw_config["skip_animations"] = any(
-        [fw_config["save_last_frame"], fw_config["from_animation_number"]]
-    )
+    fw_config["skip_animations"] = fw_config["save_last_frame"]
     fw_config["max_files_cached"] = default.getint("max_files_cached")
     if fw_config["max_files_cached"] == -1:
         fw_config["max_files_cached"] = float("inf")

manim/scene/scene.py

@@ -71,6 +71,7 @@ def __init__(self, **kwargs):
         )
         self.play_hashes_list = []
         self.mobjects = []
+        self.original_skipping_status = file_writer_config["skip_animations"]
         # TODO, remove need for foreground mobjects
         self.foreground_mobjects = []
         self.num_plays = 0
@@ -88,7 +89,7 @@ def __init__(self, **kwargs):
         self.tear_down()
         # We have to reset these settings in case of multiple renders.
         file_writer_config["skip_animations"] = False
-        self.original_skipping_status = file_writer_config["skip_animations"]
+
         self.file_writer.finish()
         self.print_end_message()
 
@@ -785,10 +786,10 @@ def update_skipping_status(self):
         """
 
         if file_writer_config["from_animation_number"]:
-            if self.num_plays == file_writer_config["from_animation_number"]:
-                file_writer_config["skip_animations"] = False
+            if self.num_plays < file_writer_config["from_animation_number"]:
+                file_writer_config["skip_animations"] = True
         if file_writer_config["upto_animation_number"]:
-            if self.num_plays >= file_writer_config["upto_animation_number"]:
+            if self.num_plays > file_writer_config["upto_animation_number"]:
                 file_writer_config["skip_animations"] = True
                 raise EndSceneEarlyException()
 
@@ -806,8 +807,13 @@ def handle_caching_play(func):
 
         def wrapper(self, *args, **kwargs):
             self.revert_to_original_skipping_status()
+            self.update_skipping_status()
             animations = self.compile_play_args_to_animation_list(*args, **kwargs)
             self.add_mobjects_from_animations(animations)
+            if file_writer_config["skip_animations"]:
+                logger.debug(f"Skipping animation {self.num_plays}")
+                func(self, *args, **kwargs)
+                return
             if not file_writer_config["disable_caching"]:
                 mobjects_on_scene = self.get_mobjects()
                 hash_play = get_hash_from_play_call(
@@ -841,6 +847,7 @@ def handle_caching_wait(func):
 
         def wrapper(self, duration=DEFAULT_WAIT_TIME, stop_condition=None):
             self.revert_to_original_skipping_status()
+            self.update_skipping_status()
             if not file_writer_config["disable_caching"]:
                 hash_wait = get_hash_from_wait_call(
                     self, self.camera, duration, stop_condition, self.get_mobjects()
@@ -880,7 +887,6 @@ def handle_play_like_call(func):
         """
 
         def wrapper(self, *args, **kwargs):
-            self.update_skipping_status()
             allow_write = not file_writer_config["skip_animations"]
             self.file_writer.begin_animation(allow_write)
             func(self, *args, **kwargs)

manim/scene/scene_file_writer.py

@@ -44,6 +44,7 @@ def __init__(self, scene, **kwargs):
         self.init_output_directories()
         self.init_audio()
         self.frame_count = 0
+        self.index_partial_movie_file = 0
 
     # Output directories and files
     def init_output_directories(self):
@@ -196,10 +197,11 @@ def get_next_partial_movie_path(self):
         result = os.path.join(
             self.partial_movie_directory,
             "{}{}".format(
-                self.scene.play_hashes_list[self.scene.num_plays],
+                self.scene.play_hashes_list[self.index_partial_movie_file],
                 file_writer_config["movie_file_extension"],
             ),
         )
+        self.index_partial_movie_file += 1
         return result
 
     def get_movie_file_path(self):

tests/control_data/videos_data/SceneWithMultipleCallsWithNFlag.json

@@ -0,0 +1,11 @@
+{
+    "name": "SceneWithMultipleCallsWithNFlag",
+    "config": {
+        "codec_name": "h264",
+        "width": 1920,
+        "height": 1080,
+        "avg_frame_rate": "60/1",
+        "duration": "4.000000",
+        "nb_frames": "240"
+    }
+}

tests/helpers/video_utils.py

@@ -4,7 +4,7 @@
 import subprocess
 import json
 
-from manim.logger import logger
+from manim.config.logger import logger
 
 
 def capture(command):

tests/test_scene_rendering/simple_scenes.py

@@ -6,3 +6,12 @@ def construct(self):
         square = Square()
         circle = Circle()
         self.play(Transform(square, circle))
+
+
+class SceneWithMultipleCalls(Scene):
+    def construct(self):
+        number = Integer(0)
+        self.add(number)
+        for i in range(10):
+            number.become(Integer(i))
+            self.play(Animation(number))

tests/test_scene_rendering/test_cli_flags.py

@@ -44,3 +44,24 @@ def test_basic_scene_l_flag(tmp_path, manim_cfg_file, simple_scenes_path):
     ]
     out, err, exit_code = capture(command)
     assert exit_code == 0, err
+
+
+@pytest.mark.slow
+@video_comparison(
+    "SceneWithMultipleCallsWithNFlag.json",
+    "videos/simple_scenes/1080p60/SceneWithMultipleCalls.mp4",
+)
+def test_n_flag(tmp_path, simple_scenes_path):
+    scene_name = "SceneWithMultipleCalls"
+    command = [
+        "python",
+        "-m",
+        "manim",
+        simple_scenes_path,
+        scene_name,
+        "-n 3,6",
+        "--media_dir",
+        str(tmp_path),
+    ]
+    _, err, exit_code = capture(command)
+    assert exit_code == 0, err