Migrate from os.path to pathlib in Testing Scripts (#2685)

* updated styling

* final commit

* fixed style

* removed exist_ok=true

* added parents=True

* potentially .exists() is the problem

* fixed style'

* fixed style on revisions

* style check processed

* Update tests/helpers/graphical_units.py

Co-authored-by: ad_chaos <[email protected]>

* fixed changes

* made get_dir_layout also accept path.

* removed small auto import error

Co-authored-by: ad_chaos <[email protected]>
Co-authored-by: Naveen M K <[email protected]>

Author: 3 people

manim/utils/commands.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
 import json
-import os
+from pathlib import Path
 from subprocess import run
-from typing import Any
+from typing import Generator
 
 __all__ = [
     "capture",
@@ -18,7 +18,7 @@ def capture(command, cwd=None, command_input=None):
     return out, err, p.returncode
 
 
-def get_video_metadata(path_to_video: str) -> dict[str, Any]:
+def get_video_metadata(path_to_video: str) -> dict[str]:
     command = [
         "ffprobe",
         "-v",
@@ -36,10 +36,10 @@ def get_video_metadata(path_to_video: str) -> dict[str, Any]:
     return json.loads(config)["streams"][0]
 
 
-def get_dir_layout(dirpath: str) -> list[str]:
+def get_dir_layout(dirpath: Path) -> Generator[str, None, None]:
     """Get list of paths relative to dirpath of all files in dir and subdirs recursively."""
-    index_files: list[str] = []
-    for root, dirs, files in os.walk(dirpath):
-        for file in files:
-            index_files.append(f"{os.path.relpath(os.path.join(root, file), dirpath)}")
-    return index_files
+    for p in dirpath.iterdir():
+        if p.is_dir():
+            yield from get_dir_layout(p)
+            continue
+        yield str(p.relative_to(dirpath))

tests/helpers/graphical_units.py

@@ -3,8 +3,8 @@
 
 from __future__ import annotations
 
-import os
 import tempfile
+from pathlib import Path
 
 import numpy as np
 
@@ -37,9 +37,9 @@ def set_test_scene(scene_object, module_name):
     config["frame_rate"] = 15
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        os.makedirs(os.path.join(tmpdir, "tex"))
-        config["text_dir"] = os.path.join(tmpdir, "text")
-        config["tex_dir"] = os.path.join(tmpdir, "tex")
+        temp_path = Path(tmpdir)
+        config["text_dir"] = temp_path / "text"
+        config["tex_dir"] = temp_path / "tex"
         scene = scene_object(skip_animations=True)
         scene.render()
         data = scene.renderer.get_frame()
@@ -48,14 +48,10 @@ def set_test_scene(scene_object, module_name):
         data == np.array([0, 0, 0, 255]),
     ), f"Control data generated for {str(scene)} only contains empty pixels."
     assert data.shape == (480, 854, 4)
-    tests_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-    path_control_data = os.path.join(
-        tests_directory,
-        "control_data",
-        "graphical_units_data",
-    )
-    path = os.path.join(path_control_data, module_name)
-    if not os.path.isdir(path):
-        os.makedirs(path)
-    np.savez_compressed(os.path.join(path, str(scene)), frame_data=data)
+    tests_directory = Path(__file__).absolute().parent.parent
+    path_control_data = Path(tests_directory) / "control_data" / "graphical_units_data"
+    path = Path(path_control_data) / module_name
+    if not path.is_dir():
+        path.mkdir(parents=True)
+    np.savez_compressed(path / str(scene), frame_data=data)
     logger.info(f"Test data for {str(scene)} saved in {path}\n")

tests/helpers/video_utils.py

@@ -3,31 +3,30 @@
 from __future__ import annotations
 
 import json
-import os
-import pathlib
+from pathlib import Path
 from typing import Any
 
 from manim import get_dir_layout, get_video_metadata, logger
 
 
-def get_section_dir_layout(dirpath: str) -> list[str]:
+def get_section_dir_layout(dirpath: Path) -> list[str]:
     """Return a list of all files in the sections directory."""
     # test if sections have been created in the first place, doesn't work with multiple scene but this isn't an issue with tests
-    if not os.path.isdir(dirpath):
+    if not dirpath.is_dir():
         return []
-    files = get_dir_layout(dirpath)
+    files = list(get_dir_layout(dirpath))
     # indicate that the sections directory has been created
     files.append(".")
     return files
 
 
-def get_section_index(metapath: str) -> list[dict[str, Any]]:
+def get_section_index(metapath: Path) -> list[dict[str, Any]]:
     """Return content of sections index file."""
-    parent_folder = pathlib.Path(metapath).parent.absolute()
+    parent_folder = metapath.parent.absolute()
     # test if sections have been created in the first place
-    if not os.path.isdir(parent_folder):
+    if not parent_folder.is_dir():
         return []
-    with open(metapath) as file:
+    with metapath.open() as file:
         index = json.load(file)
     return index
 
@@ -51,28 +50,27 @@ def save_control_data_from_video(path_to_video: str, name: str) -> None:
 
     See Also
     --------
+
     tests/utils/video_tester.py : read control data and compare with output of test
     """
-    path_to_sections = os.path.join(
-        pathlib.Path(path_to_video).parent.absolute(), "sections"
-    )
-    tests_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-    path_control_data = os.path.join(tests_directory, "control_data", "videos_data")
+    orig_path_to_sections = Path(path_to_video)
+    path_to_sections = orig_path_to_sections.parent.absolute() / "sections"
+    tests_directory = Path(__file__).absolute().parent.parent
+    path_control_data = Path(tests_directory) / "control_data" / "videos_data"
     # this is the name of the section used in the test, not the name of the test itself, it can be found as a parameter of this function
-    scene_name = "".join(os.path.basename(path_to_video).split(".")[:-1])
+    scene_name = orig_path_to_sections.stem
 
     movie_metadata = get_video_metadata(path_to_video)
     section_dir_layout = get_section_dir_layout(path_to_sections)
-    section_index = get_section_index(
-        os.path.join(path_to_sections, f"{scene_name}.json")
-    )
+    section_index = get_section_index(path_to_sections / f"{scene_name}.json")
+
     data = {
         "name": name,
         "movie_metadata": movie_metadata,
         "section_dir_layout": section_dir_layout,
         "section_index": section_index,
     }
-    path_saved = os.path.join(path_control_data, f"{name}.json")
+    path_saved = Path(path_control_data) / f"{name}.json"
     with open(path_saved, "w") as f:
         json.dump(data, f, indent=4)
     logger.info(f"Data for {name} saved in {path_saved}")

tests/test_config.py

@@ -114,26 +114,24 @@ def test_custom_dirs(tmp_path):
     ):
         scene = MyScene()
         scene.render()
+        tmp_path = Path(tmp_path)
+        assert_dir_filled(tmp_path / "test_sections")
+        assert_file_exists(tmp_path / "test_sections/MyScene.json")
 
-        assert_dir_filled(os.path.join(tmp_path, "test_sections"))
-        assert_file_exists(os.path.join(tmp_path, "test_sections", "MyScene.json"))
+        assert_dir_filled(tmp_path / "test_video")
+        assert_file_exists(tmp_path / "test_video/MyScene.mp4")
 
-        assert_dir_filled(os.path.join(tmp_path, "test_video"))
-        assert_file_exists(os.path.join(tmp_path, "test_video", "MyScene.mp4"))
-
-        assert_dir_filled(os.path.join(tmp_path, "test_partial_movie_dir"))
+        assert_dir_filled(tmp_path / "test_partial_movie_dir")
         assert_file_exists(
-            os.path.join(
-                tmp_path, "test_partial_movie_dir", "partial_movie_file_list.txt"
-            )
+            tmp_path / "test_partial_movie_dir/partial_movie_file_list.txt"
         )
 
         # TODO: another example with image output would be nice
-        assert_dir_exists(os.path.join(tmp_path, "test_images"))
+        assert_dir_exists(tmp_path / "test_images")
 
-        assert_dir_filled(os.path.join(tmp_path, "test_text"))
-        assert_dir_filled(os.path.join(tmp_path, "test_tex"))
-        assert_dir_filled(os.path.join(tmp_path, "test_log"))
+        assert_dir_filled(tmp_path / "test_text")
+        assert_dir_filled(tmp_path / "test_tex")
+        assert_dir_filled(tmp_path / "test_log")
 
 
 def test_frame_size(tmp_path):

tests/utils/GraphicalUnitTester.py

@@ -2,6 +2,7 @@
 
 import logging
 import os
+from pathlib import Path
 
 import numpy as np
 
@@ -34,27 +35,29 @@ class GraphicalUnitTester:
     def __init__(self, scene_class, module_tested, tmpdir, rgb_atol=0):
         # Disable the the logs, (--quiet is broken) TODO
         logging.disable(logging.CRITICAL)
-        tests_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-        self.path_tests_medias_cache = os.path.join(
-            tmpdir,
-            "test_graphical_units",
-            "tests_cache",
-            module_tested,
-            scene_class.__name__,
+        tests_directory = Path(__file__).absolute().parent.parent
+
+        self.path_tests_medias_cache = (
+            Path(tmpdir)
+            / "test_graphical_units"
+            / "tests_cache"
+            / module_tested
+            / scene_class.__name__
         )
-        self.path_control_data = os.path.join(
-            tests_directory,
-            "control_data",
-            "graphical_units_data",
-            module_tested,
+
+        self.path_control_data = (
+            Path(tests_directory)
+            / "control_data"
+            / "graphical_units_data"
+            / module_tested
         )
         self.rgb_atol = rgb_atol
 
         # IMPORTANT NOTE : The graphical units tests don't use for now any
         # custom manim.cfg, since it is impossible to manually select a
         # manim.cfg from a python file. (see issue #293)
-        config["text_dir"] = os.path.join(self.path_tests_medias_cache, "Text")
-        config["tex_dir"] = os.path.join(self.path_tests_medias_cache, "Tex")
+        config["text_dir"] = Path(self.path_tests_medias_cache) / "Text"
+        config["tex_dir"] = Path(self.path_tests_medias_cache) / "Tex"
 
         config["disable_caching"] = True
         config["quality"] = "low_quality"
@@ -64,7 +67,7 @@ def __init__(self, scene_class, module_tested, tmpdir, rgb_atol=0):
             config["text_dir"],
             config["tex_dir"],
         ]:
-            os.makedirs(dir_temp)
+            Path(dir_temp).mkdir(parents=True)
 
         with tempconfig({"dry_run": True}):
             if config["renderer"] == "opengl":
@@ -81,9 +84,7 @@ def _load_data(self):
         :class:`numpy.array`
             The pre-rendered frame.
         """
-        frame_data_path = os.path.join(
-            os.path.join(self.path_control_data, f"{self.scene}.npz"),
-        )
+        frame_data_path = Path(self.path_control_data) / f"{self.scene}.npz"
         return np.load(frame_data_path)["frame_data"]
 
     def _show_diff_helper(self, frame_data, expected_frame_data):

tests/utils/video_tester.py

@@ -2,8 +2,8 @@
 
 import json
 import os
-import pathlib
 from functools import wraps
+from pathlib import Path
 
 from manim import get_video_metadata
 
@@ -25,10 +25,7 @@ def check_video_data(path_control_data, path_video_gen):
         meta -> metadata
     """
     # movie file specification
-    path_sec_gen = os.path.join(
-        pathlib.Path(path_video_gen).parent.absolute(),
-        "sections",
-    )
+    path_sec_gen = Path(path_video_gen).parent.absolute() / "sections"
     control_data = load_control_data(path_control_data)
     movie_meta_gen = get_video_metadata(path_video_gen)
     movie_meta_exp = control_data["movie_metadata"]
@@ -51,8 +48,8 @@ def check_video_data(path_control_data, path_video_gen):
         raise AssertionError(f"Sections don't match:\n{mismatch}")
 
     # sections index file
-    scene_name = "".join(os.path.basename(path_video_gen).split(".")[:-1])
-    path_sec_index_gen = os.path.join(path_sec_gen, f"{scene_name}.json")
+    scene_name = Path(path_video_gen).stem
+    path_sec_index_gen = path_sec_gen / f"{scene_name}.json"
     sec_index_gen = get_section_index(path_sec_index_gen)
     sec_index_exp = control_data["section_index"]